nominatim/tools/country_info.py

   1 """
   2 Functions for importing and managing static country information.
   3 """
   4 import psycopg2.extras
   5 import yaml
   6
   7 from nominatim.db import utils as db_utils
   8 from nominatim.db.connection import connect
   9
  10 class _CountryInfo:
  11     """ Caches country-specific properties from the configuration file.
  12     """
  13
  14     def __init__(self):
  15         self._info = {}
  16
  17     def load(self, configfile):
  18         """ Load the country properties from the configuration files,
  19             if they are not loaded yet.
  20         """
  21         if not self._info:
  22             self._info = yaml.safe_load(configfile.read_text(encoding='utf-8'))
  23
  24     def items(self):
  25         """ Return tuples of (country_code, property dict) as iterable.
  26         """
  27         return self._info.items()
  28
  29
  30 _COUNTRY_INFO = _CountryInfo()
  31
  32 def setup_country_config(configfile):
  33     """ Load country properties from the configuration file.
  34         Needs to be called before using any other functions in this
  35         file.
  36     """
  37     _COUNTRY_INFO.load(configfile)
  38
  39
  40 def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
  41     """ Create and populate the tables with basic static data that provides
  42         the background for geocoding. Data is assumed to not yet exist.
  43     """
  44     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
  45     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
  46
  47     params = []
  48     for ccode, props in _COUNTRY_INFO.items():
  49         if ccode is not None and props is not None:
  50             if ignore_partitions:
  51                 partition = 0
  52             else:
  53                 partition = props.get('partition')
  54             if ',' in (props.get('languages', ',') or ','):
  55                 lang = None
  56             else:
  57                 lang = props['languages']
  58             params.append((ccode, partition, lang))
  59
  60     with connect(dsn) as conn:
  61         with conn.cursor() as cur:
  62             cur.execute_values(
  63                 """ UPDATE country_name
  64                     SET partition = part, country_default_language_code = lang
  65                     FROM (VALUES %s) AS v (cc, part, lang)
  66                     WHERE country_code = v.cc""", params)
  67         conn.commit()
  68
  69
  70 def create_country_names(conn, tokenizer, languages=None):
  71     """ Add default country names to search index. `languages` is a comma-
  72         separated list of language codes as used in OSM. If `languages` is not
  73         empty then only name translations for the given languages are added
  74         to the index.
  75     """
  76     if languages:
  77         languages = languages.split(',')
  78
  79     def _include_key(key):
  80         return key == 'name' or \
  81                (key.startswith('name:') and (not languages or key[5:] in languages))
  82
  83     with conn.cursor() as cur:
  84         psycopg2.extras.register_hstore(cur)
  85         cur.execute("""SELECT country_code, name FROM country_name
  86                        WHERE country_code is not null""")
  87
  88         with tokenizer.name_analyzer() as analyzer:
  89             for code, name in cur:
  90                 names = {'countrycode': code}
  91                 if code == 'gb':
  92                     names['short_name'] = 'UK'
  93                 if code == 'us':
  94                     names['short_name'] = 'United States'
  95
  96                 # country names (only in languages as provided)
  97                 if name:
  98                     names.update(((k, v) for k, v in name.items() if _include_key(k)))
  99
 100                 analyzer.add_country_names(code, names)
 101
 102     conn.commit()