src/nominatim_db/data/country_info.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2024 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, overload
  11 from pathlib import Path
  12 import psycopg2.extras
  13
  14 from nominatim_core.db import utils as db_utils
  15 from nominatim_core.db.connection import connect, Connection
  16 from nominatim_core.errors import UsageError
  17 from nominatim_core.config import Configuration
  18 from ..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 @overload
  91 def iterate() -> Iterable[Tuple[str, Dict[str, Any]]]:
  92     ...
  93
  94 @overload
  95 def iterate(prop: str) -> Iterable[Tuple[str, Any]]:
  96     ...
  97
  98 def iterate(prop: Optional[str] = None) -> Iterable[Tuple[str, Dict[str, Any]]]:
  99     """ Iterate over country code and properties.
 100
 101         When `prop` is None, all countries are returned with their complete
 102         set of properties.
 103
 104         If `prop` is given, then only countries are returned where the
 105         given property is set. The second item of the tuple contains only
 106         the content of the given property.
 107     """
 108     if prop is None:
 109         return _COUNTRY_INFO.items()
 110
 111     return ((c, p[prop]) for c, p in _COUNTRY_INFO.items() if prop in p)
 112
 113
 114 def setup_country_tables(dsn: str, sql_dir: Path, ignore_partitions: bool = False) -> None:
 115     """ Create and populate the tables with basic static data that provides
 116         the background for geocoding. Data is assumed to not yet exist.
 117     """
 118     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 119
 120     params = []
 121     for ccode, props in _COUNTRY_INFO.items():
 122         if ccode is not None and props is not None:
 123             if ignore_partitions:
 124                 partition = 0
 125             else:
 126                 partition = props.get('partition', 0)
 127             lang = props['languages'][0] if len(
 128                 props['languages']) == 1 else None
 129
 130             params.append((ccode, props['names'], lang, partition))
 131     with connect(dsn) as conn:
 132         with conn.cursor() as cur:
 133             psycopg2.extras.register_hstore(cur)
 134             cur.execute(
 135                 """ CREATE TABLE public.country_name (
 136                         country_code character varying(2),
 137                         name public.hstore,
 138                         derived_name public.hstore,
 139                         country_default_language_code text,
 140                         partition integer
 141                     ); """)
 142             cur.execute_values(
 143                 """ INSERT INTO public.country_name
 144                     (country_code, name, country_default_language_code, partition) VALUES %s
 145                 """, params)
 146         conn.commit()
 147
 148
 149 def create_country_names(conn: Connection, tokenizer: AbstractTokenizer,
 150                          languages: Optional[Container[str]] = None) -> None:
 151     """ Add default country names to search index. `languages` is a comma-
 152         separated list of language codes as used in OSM. If `languages` is not
 153         empty then only name translations for the given languages are added
 154         to the index.
 155     """
 156     def _include_key(key: str) -> bool:
 157         return ':' not in key or not languages or \
 158                key[key.index(':') + 1:] in languages
 159
 160     with conn.cursor() as cur:
 161         psycopg2.extras.register_hstore(cur)
 162         cur.execute("""SELECT country_code, name FROM country_name
 163                        WHERE country_code is not null""")
 164
 165         with tokenizer.name_analyzer() as analyzer:
 166             for code, name in cur:
 167                 names = {'countrycode': code}
 168
 169                 # country names (only in languages as provided)
 170                 if name:
 171                     names.update({k : v for k, v in name.items() if _include_key(k)})
 172
 173                 analyzer.add_country_names(code, names)
 174
 175     conn.commit()