clean up

[nominatim.git] / nominatim / tools / country_info.py

diff --git a/nominatim/tools/country_info.py b/nominatim/tools/country_info.py
index 1f7b799610743926f513f67fe3e40b20e76e6397..9ed666ff21c61acf13bf09633e437ffba7818c4e 100644 (file)
--- a/nominatim/tools/country_info.py
+++ b/nominatim/tools/country_info.py
@@ -1,8 +1,14 @@
+# SPDX-License-Identifier: GPL-2.0-only
+#
+# This file is part of Nominatim. (https://nominatim.org)
+#
+# Copyright (C) 2022 by the Nominatim developer community.
+# For a full list of authors see the git log.

 """
 Functions for importing and managing static country information.
 """
 """
 Functions for importing and managing static country information.
 """

+import json

 import psycopg2.extras
 import psycopg2.extras

-import yaml

 
 from nominatim.db import utils as db_utils
 from nominatim.db.connection import connect
 
 from nominatim.db import utils as db_utils
 from nominatim.db.connection import connect


@@ -14,32 +20,61 @@ class _CountryInfo:
     def __init__(self):
         self._info = {}
 
     def __init__(self):
         self._info = {}
 

-    def load(self, configfile):
+
+    def load(self, config):
+        """ Load the country properties from the configuration files,
+            if they are not loaded yet.
+        """

         if not self._info:
         if not self._info:

-            self._info = yaml.safe_load(configfile.read_text())
+            self._info = config.load_sub_configuration('country_settings.yaml')
+            self._key_prefix = 'name'
+            # Convert languages into a list for simpler handling.
+            for prop in self._info.values():
+                if 'languages' not in prop:
+                    prop['languages'] = []
+                elif not isinstance(prop['languages'], list):
+                    prop['languages'] = [x.strip()
+                                         for x in prop['languages'].split(',')]
+                if 'names' not in prop:
+                    prop['names'][self._key_prefix] = {}

 
     def items(self):
 
     def items(self):

+        """ Return tuples of (country_code, property dict) as iterable.
+        """

         return self._info.items()
 
         return self._info.items()
 

+    def key_prefix(self):
+        """ Return the prefix that will be attached to the keys of the country 
+            names values when storing them in the database
+        """
+        return self._key_prefix
+

 
 _COUNTRY_INFO = _CountryInfo()
 
 
 _COUNTRY_INFO = _CountryInfo()
 

-def setup_country_config(configfile):
+def setup_country_config(config):

     """ Load country properties from the configuration file.
         Needs to be called before using any other functions in this
         file.
     """
     """ Load country properties from the configuration file.
         Needs to be called before using any other functions in this
         file.
     """

-    _COUNTRY_INFO.load(configfile)
-    print(_COUNTRY_INFO._info)
+    _COUNTRY_INFO.load(config)
+
+
+def iterate():
+    """ Iterate over country code and properties.
+    """
+    return _COUNTRY_INFO.items()

 
 
 def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
     """ Create and populate the tables with basic static data that provides
         the background for geocoding. Data is assumed to not yet exist.
     """
 
 
 def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
     """ Create and populate the tables with basic static data that provides
         the background for geocoding. Data is assumed to not yet exist.
     """

-    db_utils.execute_file(dsn, sql_dir / 'country_name.sql')

     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 
     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 

+    def add_prefix_to_keys(names, prefix):
+        return {prefix+':'+k: v for k, v in names.items()}
+

     params = []
     for ccode, props in _COUNTRY_INFO.items():
         if ccode is not None and props is not None:
     params = []
     for ccode, props in _COUNTRY_INFO.items():
         if ccode is not None and props is not None:


@@ -47,19 +82,24 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
                 partition = 0
             else:
                 partition = props.get('partition')
                 partition = 0
             else:
                 partition = props.get('partition')

-            if ',' in (props.get('languages', ',') or ','):
-                lang = None
-            else:
-                lang = props['languages']
-            params.append((ccode, partition, lang))
-
+            lang = props['languages'][0] if len(props['languages']) == 1 else None
+            name = add_prefix_to_keys(props.get('names').get(_COUNTRY_INFO.key_prefix()), _COUNTRY_INFO.key_prefix())
+            name = json.dumps(name, ensure_ascii=False, separators=(', ', '=>'))
+            params.append((ccode, name[1:-1], lang, partition))

     with connect(dsn) as conn:
         with conn.cursor() as cur:
     with connect(dsn) as conn:
         with conn.cursor() as cur:

+            cur.execute(
+                """ CREATE TABLE public.country_name (
+                        country_code character varying(2),
+                        name public.hstore,
+                        derived_name public.hstore,
+                        country_default_language_code text,
+                        partition integer
+                    ); """)

             cur.execute_values(
             cur.execute_values(

-                """ UPDATE country_name
-                    SET partition = part, country_default_language_code = lang
-                    FROM (VALUES %s) AS v (cc, part, lang)
-                    WHERE country_code = v.cc""", params)
+                """ INSERT INTO public.country_name
+                    (country_code, name, country_default_language_code, partition) VALUES %s
+                """, params)

         conn.commit()
 
 
         conn.commit()
 
 


@@ -73,8 +113,8 @@ def create_country_names(conn, tokenizer, languages=None):
         languages = languages.split(',')
 
     def _include_key(key):
         languages = languages.split(',')
 
     def _include_key(key):

-        return key == 'name' or \
-               (key.startswith('name:') and (not languages or key[5:] in languages))
+        return key == _COUNTRY_INFO.key_prefix() or \
+               (key.startswith(_COUNTRY_INFO.key_prefix()+':') and (not languages or key[len(_COUNTRY_INFO.key_prefix())+1:] in languages))

 
     with conn.cursor() as cur:
         psycopg2.extras.register_hstore(cur)
 
     with conn.cursor() as cur:
         psycopg2.extras.register_hstore(cur)