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