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