]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/data/country_info.py
add typing information for postcode formatter
[nominatim.git] / 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, overload
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 @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()