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