X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/6be2077d922098d65f069557dea55542c15a5c98..baee6f3de09226c3dc41cb2314a0ac348e865561:/nominatim/tools/country_info.py?ds=sidebyside diff --git a/nominatim/tools/country_info.py b/nominatim/tools/country_info.py index 87be7290..d754b4dd 100644 --- a/nominatim/tools/country_info.py +++ b/nominatim/tools/country_info.py @@ -7,12 +7,35 @@ """ Functions for importing and managing static country information. """ -import json -from io import StringIO import psycopg2.extras from nominatim.db import utils as db_utils from nominatim.db.connection import connect +from nominatim.errors import UsageError + +def _flatten_name_list(names): + if names is None: + return {} + + if not isinstance(names, dict): + raise UsageError("Expected key-value list for names in country_settings.py") + + flat = {} + for prefix, remain in names.items(): + if isinstance(remain, str): + flat[prefix] = remain + elif not isinstance(remain, dict): + raise UsageError("Entries in names must be key-value lists.") + else: + for suffix, name in remain.items(): + if suffix == 'default': + flat[prefix] = name + else: + flat[f'{prefix}:{suffix}'] = name + + return flat + + class _CountryInfo: """ Caches country-specific properties from the configuration file. @@ -28,24 +51,31 @@ class _CountryInfo: """ if not self._info: self._info = config.load_sub_configuration('country_settings.yaml') - # Convert languages into a list for simpler handling. for prop in self._info.values(): + # Convert languages into a list for simpler handling. if 'languages' not in prop: prop['languages'] = [] elif not isinstance(prop['languages'], list): prop['languages'] = [x.strip() for x in prop['languages'].split(',')] - if 'name' not in prop: - prop['name'] = {} + prop['names'] = _flatten_name_list(prop.get('names')) + def items(self): """ Return tuples of (country_code, property dict) as iterable. """ return self._info.items() + def get(self, country_code): + """ Get country information for the country with the given country code. + """ + return self._info.get(country_code, {}) + + _COUNTRY_INFO = _CountryInfo() + def setup_country_config(config): """ Load country properties from the configuration file. Needs to be called before using any other functions in this @@ -54,10 +84,20 @@ def setup_country_config(config): _COUNTRY_INFO.load(config) -def iterate(): +def iterate(prop=None): """ Iterate over country code and properties. + + When `prop` is None, all countries are returned with their complete + set of properties. + + If `prop` is given, then only countries are returned where the + given property is set. The second item of the tuple contains only + the content of the given property. """ - return _COUNTRY_INFO.items() + if prop is None: + return _COUNTRY_INFO.items() + + return ((c, p[prop]) for c, p in _COUNTRY_INFO.items() if prop in p) def setup_country_tables(dsn, sql_dir, ignore_partitions=False): @@ -66,22 +106,20 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False): """ db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz') - params, country_names_data = [], '' + params = [] for ccode, props in _COUNTRY_INFO.items(): if ccode is not None and props is not None: if ignore_partitions: partition = 0 else: partition = props.get('partition') - lang = props['languages'][0] if len(props['languages']) == 1 else None - params.append((ccode, partition, lang)) - - name = json.dumps(props.get('names').get('name'), ensure_ascii=False, - separators=(', ', '=>')) - country_names_data += ccode + '\t' + name[1:-1] + '\n' + lang = props['languages'][0] if len( + props['languages']) == 1 else None + params.append((ccode, props['names'], lang, partition)) with connect(dsn) as conn: with conn.cursor() as cur: + psycopg2.extras.register_hstore(cur) cur.execute( """ CREATE TABLE public.country_name ( country_code character varying(2), @@ -90,13 +128,10 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False): country_default_language_code text, partition integer ); """) - data = StringIO(country_names_data) - cur.copy_from(data, 'country_name', columns=('country_code', 'name')) cur.execute_values( - """ UPDATE country_name - SET partition = part, country_default_language_code = lang - FROM (VALUES %s) AS v (cc, part, lang) - WHERE country_code = v.cc""", params) + """ INSERT INTO public.country_name + (country_code, name, country_default_language_code, partition) VALUES %s + """, params) conn.commit() @@ -106,12 +141,9 @@ def create_country_names(conn, tokenizer, languages=None): empty then only name translations for the given languages are added to the index. """ - if languages: - languages = languages.split(',') - def _include_key(key): - return key == 'default' or \ - (not languages or key in languages) + return ':' not in key or not languages or \ + key[key.index(':') + 1:] in languages with conn.cursor() as cur: psycopg2.extras.register_hstore(cur) @@ -121,14 +153,10 @@ def create_country_names(conn, tokenizer, languages=None): with tokenizer.name_analyzer() as analyzer: for code, name in cur: names = {'countrycode': code} - if code == 'gb': - names['short_name'] = 'UK' - if code == 'us': - names['short_name'] = 'United States' # country names (only in languages as provided) if name: - names.update(((k, v) for k, v in name.items() if _include_key(k))) + names.update({k : v for k, v in name.items() if _include_key(k)}) analyzer.add_country_names(code, names)