From d0c44431d0e1386f0d3278f6a8bf84a01381e36a Mon Sep 17 00:00:00 2001 From: Sarah Hoffmann Date: Thu, 7 Jul 2022 17:31:20 +0200 Subject: [PATCH] add typing information for place_info and country_info --- nominatim/data/country_info.py | 31 ++++++++++++++++++------------- nominatim/data/place_info.py | 17 +++++++++-------- nominatim/tokenizer/base.py | 4 ++-- 3 files changed, 29 insertions(+), 23 deletions(-) diff --git a/nominatim/data/country_info.py b/nominatim/data/country_info.py index d754b4dd..ada76325 100644 --- a/nominatim/data/country_info.py +++ b/nominatim/data/country_info.py @@ -7,13 +7,17 @@ """ Functions for importing and managing static country information. """ +from typing import Dict, Any, Iterable, Tuple, Optional, Container +from pathlib import Path import psycopg2.extras from nominatim.db import utils as db_utils -from nominatim.db.connection import connect +from nominatim.db.connection import connect, Connection from nominatim.errors import UsageError +from nominatim.config import Configuration +from nominatim.tokenizer.base import AbstractTokenizer -def _flatten_name_list(names): +def _flatten_name_list(names: Any) -> Dict[str, str]: if names is None: return {} @@ -41,11 +45,11 @@ class _CountryInfo: """ Caches country-specific properties from the configuration file. """ - def __init__(self): - self._info = {} + def __init__(self) -> None: + self._info: Dict[str, Dict[str, Any]] = {} - def load(self, config): + def load(self, config: Configuration) -> None: """ Load the country properties from the configuration files, if they are not loaded yet. """ @@ -61,12 +65,12 @@ class _CountryInfo: prop['names'] = _flatten_name_list(prop.get('names')) - def items(self): + def items(self) -> Iterable[Tuple[str, Dict[str, Any]]]: """ Return tuples of (country_code, property dict) as iterable. """ return self._info.items() - def get(self, country_code): + def get(self, country_code: str) -> Dict[str, Any]: """ Get country information for the country with the given country code. """ return self._info.get(country_code, {}) @@ -76,7 +80,7 @@ class _CountryInfo: _COUNTRY_INFO = _CountryInfo() -def setup_country_config(config): +def setup_country_config(config: Configuration) -> None: """ Load country properties from the configuration file. Needs to be called before using any other functions in this file. @@ -84,7 +88,7 @@ def setup_country_config(config): _COUNTRY_INFO.load(config) -def iterate(prop=None): +def iterate(prop: Optional[str] = None) -> Iterable[Tuple[str, Dict[str, Any]]]: """ Iterate over country code and properties. When `prop` is None, all countries are returned with their complete @@ -100,7 +104,7 @@ def iterate(prop=None): 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): +def setup_country_tables(dsn: str, sql_dir: Path, ignore_partitions: bool = False) -> None: """ Create and populate the tables with basic static data that provides the background for geocoding. Data is assumed to not yet exist. """ @@ -112,7 +116,7 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False): if ignore_partitions: partition = 0 else: - partition = props.get('partition') + partition = props.get('partition', 0) lang = props['languages'][0] if len( props['languages']) == 1 else None @@ -135,13 +139,14 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False): conn.commit() -def create_country_names(conn, tokenizer, languages=None): +def create_country_names(conn: Connection, tokenizer: AbstractTokenizer, + languages: Optional[Container[str]] = None) -> None: """ Add default country names to search index. `languages` is a comma- separated list of language codes as used in OSM. If `languages` is not empty then only name translations for the given languages are added to the index. """ - def _include_key(key): + def _include_key(key: str) -> bool: return ':' not in key or not languages or \ key[key.index(':') + 1:] in languages diff --git a/nominatim/data/place_info.py b/nominatim/data/place_info.py index d2ba3979..96912a61 100644 --- a/nominatim/data/place_info.py +++ b/nominatim/data/place_info.py @@ -8,18 +8,19 @@ Wrapper around place information the indexer gets from the database and hands to the tokenizer. """ +from typing import Optional, Mapping, Any class PlaceInfo: """ Data class containing all information the tokenizer gets about a place it should process the names for. """ - def __init__(self, info): + def __init__(self, info: Mapping[str, Any]) -> None: self._info = info @property - def name(self): + def name(self) -> Optional[Mapping[str, str]]: """ A dictionary with the names of the place or None if the place has no names. """ @@ -27,7 +28,7 @@ class PlaceInfo: @property - def address(self): + def address(self) -> Optional[Mapping[str, str]]: """ A dictionary with the address elements of the place or None if no address information is available. """ @@ -35,7 +36,7 @@ class PlaceInfo: @property - def country_code(self): + def country_code(self) -> Optional[str]: """ The country code of the country the place is in. Guaranteed to be a two-letter lower-case string or None, if no country could be found. @@ -44,20 +45,20 @@ class PlaceInfo: @property - def rank_address(self): + def rank_address(self) -> int: """ The computed rank address before rank correction. """ - return self._info.get('rank_address') + return self._info.get('rank_address', 0) - def is_a(self, key, value): + def is_a(self, key: str, value: str) -> bool: """ Check if the place's primary tag corresponds to the given key and value. """ return self._info.get('class') == key and self._info.get('type') == value - def is_country(self): + def is_country(self) -> bool: """ Check if the place is a valid country boundary. """ return self.rank_address == 4 \ diff --git a/nominatim/tokenizer/base.py b/nominatim/tokenizer/base.py index 70a54bfd..5a3d3b12 100644 --- a/nominatim/tokenizer/base.py +++ b/nominatim/tokenizer/base.py @@ -28,7 +28,7 @@ class AbstractAnalyzer(ABC): return self - def __exit__(self, exc_type, exc_value, traceback) -> None: + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: self.close() @@ -95,7 +95,7 @@ class AbstractAnalyzer(ABC): @abstractmethod - def add_country_names(self, country_code: str, names: Dict[str, str]): + def add_country_names(self, country_code: str, names: Dict[str, str]) -> None: """ Add the given names to the tokenizer's list of country tokens. Arguments: -- 2.39.5