Load and update pycountry data from script

issue8346
review285501002

CHANGELOG

@@ -1,3 +1,5 @@
+* Replace XML data by import script
+
 Version 5.2.0 - 2019-05-06
 * Bug fixes (see mercurial logs for details)
 * Remove starting wildcard when searching on codes and numbers

country.py

@@ -1,7 +1,9 @@
 # This file is part of Tryton.  The COPYRIGHT file at the top level of
 # this repository contains the full copyright notices and license terms.
 from trytond.model import ModelView, ModelSQL, DeactivableMixin, fields
+from trytond.pool import Pool
 from trytond.pyson import Eval
+from trytond.transaction import Transaction
 from trytond.tools import lstrip_wildcard
 
 __all__ = ['Country', 'Subdivision', 'Zip']
@@ -28,6 +30,11 @@ class Country(DeactivableMixin, ModelSQL, ModelView):
 
     @classmethod
     def __register__(cls, module_name):
+        pool = Pool()
+        Data = pool.get('ir.model.data')
+        data = Data.__table__()
+        cursor = Transaction().connection.cursor()
+
         super(Country, cls).__register__(module_name)
 
         table = cls.__table_handler__(module_name)
@@ -39,6 +46,10 @@ class Country(DeactivableMixin, ModelSQL, ModelView):
         # Migration from 3.8: remove required on code
         table.not_null_action('code', 'remove')
 
+        # Migration from 5.2: remove country data
+        cursor.execute(*data.delete(where=(data.module == 'country')
+                & (data.model == cls.__name__)))
+
     @classmethod
     def search_rec_name(cls, name, clause):
         code_value = clause[2]
@@ -201,6 +212,19 @@ class Subdivision(DeactivableMixin, ModelSQL, ModelView):
         super(Subdivision, cls).__setup__()
         cls._order.insert(0, ('code', 'ASC'))
 
+    @classmethod
+    def __register__(cls, module_name):
+        pool = Pool()
+        Data = pool.get('ir.model.data')
+        data = Data.__table__()
+        cursor = Transaction().connection.cursor()
+
+        super().__register__(module_name)
+
+        # Migration from 5.2: remove country data
+        cursor.execute(*data.delete(where=(data.module == 'country')
+                & (data.model == cls.__name__)))
+
     @classmethod
     def search_rec_name(cls, name, clause):
         return ['OR',

doc/index.rst

@@ -24,3 +24,13 @@ Zip
 ***
 
 A zip links a zip code of a country to a city.
+
+Scripts
+*******
+
+There are two scripts:
+
+* `trytond_import_countries` to create and update countries and
+  subdivisions from the ISO database.
+
+* `trytond_import_zip` to create zips by country from the geonames database.

scripts/countries.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python
-# This file is part of Tryton.  The COPYRIGHT file at the top level of
-# this repository contains the full copyright notices and license terms.
-
-import sys
-import pycountry
-
-sys.stdout.write('<?xml version="1.0"?>\n')
-sys.stdout.write('<tryton>\n')
-sys.stdout.write('    <data skiptest="1" grouped="1">\n')
-
-for country in pycountry.countries:
-    record = '''
-        <record model="country.country" id="%s">
-            <field name="name">%s</field>
-            <field name="code">%s</field>
-            <field name="code3">%s</field>
-            <field name="code_numeric">%s</field>
-        </record>\n''' % (country.alpha_2.lower(), country.name,
-                country.alpha_2, country.alpha_3, country.numeric)
-    sys.stdout.write(record)
-sys.stdout.write('    </data>\n')
-
-subdivision_codes = {s.code.lower() for s in pycountry.subdivisions}
-existing_parents = set()
-while existing_parents < subdivision_codes:
-    sys.stdout.write('    <data skiptest="1" grouped="1">\n')
-    new_parents = set()
-    for subdivision in pycountry.subdivisions:
-        if (subdivision.code.lower() in existing_parents
-                or (subdivision.parent_code
-                    and subdivision.parent_code.lower() not in
-                    existing_parents)):
-            continue
-        new_parents.add(subdivision.code.lower())
-        record = '''
-        <record model="country.subdivision" id="%s">
-            <field name="name">%s</field>
-            <field name="code">%s</field>
-            <field name="type">%s</field>\n''' % (subdivision.code.lower(),
-                    subdivision.name, subdivision.code,
-                    subdivision.type.lower())
-        if subdivision.parent_code:
-            record += '''\
-            <field name="parent" ref="%s"/>\n''' % \
-                subdivision.parent.code.lower()
-        record += '''\
-            <field name="country" ref="%s"/>
-        </record>\n''' % subdivision.country.alpha_2.lower()
-
-        sys.stdout.write(record)
-    sys.stdout.write('    </data>\n')
-    existing_parents |= new_parents
-
-sys.stdout.write('</tryton>\n')

scripts/trytond_import_countries

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# This file is part of Tryton.  The COPYRIGHT file at the top level of
+# this repository contains the full copyright notices and license terms.
+import gettext
+import os
+import sys
+from argparse import ArgumentParser
+
+import pycountry
+
+try:
+    from progressbar import ProgressBar, Bar, ETA, SimpleProgress
+except ImportError:
+    ProgressBar = None
+
+try:
+    from proteus import Model, config
+except ImportError:
+    prog = os.path.basename(sys.argv[0])
+    sys.exit("proteus must be installed to use %s" % prog)
+
+
+def _progress(iterable):
+    if ProgressBar:
+        pbar = ProgressBar(
+            widgets=[SimpleProgress(), Bar(), ETA()])
+    else:
+        pbar = iter
+    return pbar(iterable)
+
+
+def _get_language_codes():
+    Language = Model.get('ir.lang')
+    languages = Language.find([('translatable', '=', True)])
+    for l in languages:
+        yield l.code
+
+
+def get_countries():
+    Country = Model.get('country.country')
+
+    with config.get_config().set_context(active_test=False):
+        return {c.code: c for c in Country.find([])}
+
+
+def update_countries(countries):
+    print("Update countries")
+    Country = Model.get('country.country')
+
+    records = []
+    for country in _progress(pycountry.countries):
+        code = country.alpha_2
+        if code in countries:
+            record = countries[code]
+        else:
+            record = Country(code=code)
+        record.name = country.name
+        record.code3 = country.alpha_3
+        record.code_numeric = country.numeric
+        records.append(record)
+
+    Country.save(records)
+    return {c.code: c for c in records}
+
+
+def translate_countries(countries):
+    Country = Model.get('country.country')
+
+    current_config = config.get_config()
+    for code in _get_language_codes():
+        try:
+            gnutranslation = gettext.translation(
+                'iso3166', pycountry.LOCALES_DIR, languages=[code])
+        except IOError:
+            continue
+        print("Update countries %s" % code)
+        with current_config.set_context(language=code):
+            records = []
+            for country in _progress(pycountry.countries):
+                record = Country(countries[country.alpha_2].id)
+                record.name = gnutranslation.gettext(country.name)
+                records.append(record)
+            Country.save(records)
+
+
+def get_subdivisions():
+    Subdivision = Model.get('country.subdivision')
+
+    with config.get_config().set_context(active_test=False):
+        return {(s.country.code, s.code): s for s in Subdivision.find([])}
+
+
+def update_subdivisions(countries, subdivisions):
+    print("Update subdivisions")
+    Subdivision = Model.get('country.subdivision')
+
+    records = []
+    for subdivision in _progress(pycountry.subdivisions):
+        code = subdivision.code
+        country_code = subdivision.country_code
+        if (country_code, code) in subdivisions:
+            record = subdivisions[(country_code, code)]
+        else:
+            record = Subdivision(code=code, country=countries[country_code])
+        record.name = subdivision.name
+        record.type = subdivision.type.lower()
+        records.append(record)
+
+    Subdivision.save(records)
+    return {(s.country.code, s.code): s for s in records}
+
+
+def update_subdivisions_parent(subdivisions):
+    print("Update subdivisions parent")
+    Subdivision = Model.get('country.subdivision')
+
+    records = []
+    for subdivision in _progress(pycountry.subdivisions):
+        code = subdivision.code
+        country_code = subdivision.country_code
+        record = subdivisions[(country_code, code)]
+        if subdivision.parent_code:
+            record.parent = subdivisions[
+                (country_code, subdivision.parent_code)]
+        else:
+            record.parent = None
+        records.append(record)
+    Subdivision.save(records)
+
+
+def translate_subdivisions(subdivisions):
+    Subdivision = Model.get('country.subdivision')
+
+    current_config = config.get_config()
+    for code in _get_language_codes():
+        try:
+            gnutranslation = gettext.translation(
+                'iso3166-2', pycountry.LOCALES_DIR, languages=[code])
+        except IOError:
+            continue
+        print("Update subdivisions %s" % code)
+        with current_config.set_context(language=code):
+            records = []
+            for subdivision in _progress(pycountry.subdivisions):
+                record = Subdivision(subdivisions[
+                        (subdivision.country_code, subdivision.code)].id)
+                record.name = gnutranslation.gettext(subdivision.name)
+                records.append(record)
+            Subdivision.save(records)
+
+
+def main(database, config_file=None):
+    config.set_trytond(database, config_file=config_file)
+
+    countries = get_countries()
+    countries = update_countries(countries)
+    translate_countries(countries)
+    subdivisions = get_subdivisions()
+    subdivisions = update_subdivisions(countries, subdivisions)
+    update_subdivisions_parent(subdivisions)
+    translate_subdivisions(subdivisions)
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('-d', '--database', dest='database')
+    parser.add_argument('-c', '--config', dest='config_file',
+        help='the trytond config file')
+
+    args = parser.parse_args()
+    if not args.database:
+        parser.error('Missing database')
+    main(args.database, args.config_file)

setup.py

@@ -79,7 +79,10 @@ setup(name=name,
         'trytond.modules.country': (info.get('xml', [])
             + ['tryton.cfg', 'view/*.xml', 'locale/*.po', 'icons/*.svg']),
         },
-    scripts=['scripts/trytond_import_zip'],
+    scripts=[
+        'scripts/trytond_import_countries',
+        'scripts/trytond_import_zip',
+        ],
     classifiers=[
         'Development Status :: 5 - Production/Stable',
         'Environment :: Plugins',
@@ -120,6 +123,7 @@ setup(name=name,
     python_requires='>=3.5',
     install_requires=requires,
     extras_require={
+        'data': ['pycountry', get_require_version('proteus')],
         'GeoNames': [get_require_version('proteus')],
         },
     dependency_links=dependency_links,

tryton.cfg

@@ -5,4 +5,3 @@ depends:
     res
 xml:
     country.xml
-    data.xml