nominatim/tools/country_info.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2022 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Functions for importing and managing static country information.
   9 """
  10 import psycopg2.extras
  11
  12 from nominatim.db import utils as db_utils
  13 from nominatim.db.connection import connect
  14 from nominatim.errors import UsageError
  15
  16 def _flatten_name_list(names):
  17     if names is None:
  18         return {}
  19
  20     if not isinstance(names, dict):
  21         raise UsageError("Expected key-value list for names in country_settings.py")
  22
  23     flat = {}
  24     for prefix, remain in names.items():
  25         if isinstance(remain, str):
  26             flat[prefix] = remain
  27         elif not isinstance(remain, dict):
  28             raise UsageError("Entries in names must be key-value lists.")
  29         else:
  30             for suffix, name in remain.items():
  31                 if suffix == 'default':
  32                     flat[prefix] = name
  33                 else:
  34                     flat[f'{prefix}:{suffix}'] = name
  35
  36     return flat
  37
  38
  39
  40 class _CountryInfo:
  41     """ Caches country-specific properties from the configuration file.
  42     """
  43
  44     def __init__(self):
  45         self._info = {}
  46
  47
  48     def load(self, config):
  49         """ Load the country properties from the configuration files,
  50             if they are not loaded yet.
  51         """
  52         if not self._info:
  53             self._info = config.load_sub_configuration('country_settings.yaml')
  54             for prop in self._info.values():
  55                 # Convert languages into a list for simpler handling.
  56                 if 'languages' not in prop:
  57                     prop['languages'] = []
  58                 elif not isinstance(prop['languages'], list):
  59                     prop['languages'] = [x.strip()
  60                                          for x in prop['languages'].split(',')]
  61                 prop['names'] = _flatten_name_list(prop.get('names'))
  62
  63
  64     def items(self):
  65         """ Return tuples of (country_code, property dict) as iterable.
  66         """
  67         return self._info.items()
  68
  69     def get(self, country_code):
  70         """ Get country information for the country with the given country code.
  71         """
  72         return self._info.get(country_code, {})
  73
  74
  75
  76 _COUNTRY_INFO = _CountryInfo()
  77
  78
  79 def setup_country_config(config):
  80     """ Load country properties from the configuration file.
  81         Needs to be called before using any other functions in this
  82         file.
  83     """
  84     _COUNTRY_INFO.load(config)
  85
  86
  87 def iterate(prop=None):
  88     """ Iterate over country code and properties.
  89
  90         When `prop` is None, all countries are returned with their complete
  91         set of properties.
  92
  93         If `prop` is given, then only countries are returned where the
  94         given property is set. The second item of the tuple contains only
  95         the content of the given property.
  96     """
  97     if prop is None:
  98         return _COUNTRY_INFO.items()
  99
 100     return ((c, p[prop]) for c, p in _COUNTRY_INFO.items() if prop in p)
 101
 102
 103 def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
 104     """ Create and populate the tables with basic static data that provides
 105         the background for geocoding. Data is assumed to not yet exist.
 106     """
 107     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 108
 109     params = []
 110     for ccode, props in _COUNTRY_INFO.items():
 111         if ccode is not None and props is not None:
 112             if ignore_partitions:
 113                 partition = 0
 114             else:
 115                 partition = props.get('partition')
 116             lang = props['languages'][0] if len(
 117                 props['languages']) == 1 else None
 118
 119             params.append((ccode, props['names'], lang, partition))
 120     with connect(dsn) as conn:
 121         with conn.cursor() as cur:
 122             psycopg2.extras.register_hstore(cur)
 123             cur.execute(
 124                 """ CREATE TABLE public.country_name (
 125                         country_code character varying(2),
 126                         name public.hstore,
 127                         derived_name public.hstore,
 128                         country_default_language_code text,
 129                         partition integer
 130                     ); """)
 131             cur.execute_values(
 132                 """ INSERT INTO public.country_name
 133                     (country_code, name, country_default_language_code, partition) VALUES %s
 134                 """, params)
 135         conn.commit()
 136
 137
 138 def create_country_names(conn, tokenizer, languages=None):
 139     """ Add default country names to search index. `languages` is a comma-
 140         separated list of language codes as used in OSM. If `languages` is not
 141         empty then only name translations for the given languages are added
 142         to the index.
 143     """
 144     def _include_key(key):
 145         return ':' not in key or not languages or \
 146                key[key.index(':') + 1:] in languages
 147
 148     with conn.cursor() as cur:
 149         psycopg2.extras.register_hstore(cur)
 150         cur.execute("""SELECT country_code, name FROM country_name
 151                        WHERE country_code is not null""")
 152
 153         with tokenizer.name_analyzer() as analyzer:
 154             for code, name in cur:
 155                 names = {'countrycode': code}
 156
 157                 # country names (only in languages as provided)
 158                 if name:
 159                     names.update({k : v for k, v in name.items() if _include_key(k)})
 160
 161                 analyzer.add_country_names(code, names)
 162
 163     conn.commit()