add type hints

geonamescache/__init__.py

@@ -7,54 +7,62 @@
 
 import json
 import os
+from typing import Any, Dict, List, Mapping, Optional, Tuple, TypeVar
 
 from . import geonamesdata
+from .types import (City, CitySearchAttribute, Continent, ContinentCode,
+                    Country, GeoNameIdStr, ISOStr, USCounty, USState,
+                    USStateCode, USStateName)
+
+TDict = TypeVar('TDict', bound=Mapping[str, Any])
 
 
 class GeonamesCache:
 
-    us_states = geonamesdata.us_states
-    continents = None
-    countries = None
-    cities = None
-    cities_items = None
-    cities_by_names = {}
-    us_counties = None
+    us_states: Dict[USStateCode, USState] = geonamesdata.us_states
+    continents: Optional[Dict[ContinentCode, Continent]] = None
+    countries: Optional[Dict[ISOStr, Country]] = None
+    cities: Optional[Dict[GeoNameIdStr, City]] = None
+    cities_items: Optional[List[Tuple[GeoNameIdStr, City]]] = None
+    cities_by_names: Dict[str, List[Dict[GeoNameIdStr, City]]] = {}
+    us_counties: Optional[List[USCounty]] = None
 
-    def __init__(self, min_city_population=15000):
+    def __init__(self, min_city_population: int = 15000):
         self.min_city_population = min_city_population
 
-    def get_dataset_by_key(self, dataset, key):
+    def get_dataset_by_key(
+        self, dataset: Dict[Any, TDict], key: str
+    ) -> Dict[Any, TDict]:
         return dict((d[key], d) for c, d in list(dataset.items()))
 
-    def get_continents(self):
+    def get_continents(self) -> Dict[ContinentCode, Continent]:
         if self.continents is None:
             self.continents = self._load_data(
                 self.continents, 'continents.json')
         return self.continents
 
-    def get_countries(self):
+    def get_countries(self) -> Dict[ISOStr, Country]:
         if self.countries is None:
             self.countries = self._load_data(self.countries, 'countries.json')
         return self.countries
 
-    def get_us_states(self):
+    def get_us_states(self) -> Dict[USStateCode, USState]:
         return self.us_states
 
-    def get_countries_by_names(self):
+    def get_countries_by_names(self) -> Dict[str, Country]:
         return self.get_dataset_by_key(self.get_countries(), 'name')
 
-    def get_us_states_by_names(self):
+    def get_us_states_by_names(self) -> Dict[USStateName, USState]:
         return self.get_dataset_by_key(self.get_us_states(), 'name')
 
-    def get_cities(self):
+    def get_cities(self) -> Dict[GeoNameIdStr, City]:
         """Get a dictionary of cities keyed by geonameid."""
 
         if self.cities is None:
             self.cities = self._load_data(self.cities, f'cities{self.min_city_population}.json')
         return self.cities
 
-    def get_cities_by_name(self, name):
+    def get_cities_by_name(self, name: str) -> List[Dict[GeoNameIdStr, City]]:
         """Get a list of city dictionaries with the given name.
 
         City names cannot be used as keys, as they are not unique.
@@ -72,7 +80,13 @@ def get_us_counties(self):
             self.us_counties = self._load_data(self.us_counties, 'us_counties.json')
         return self.us_counties
 
-    def search_cities(self, query, attribute='alternatenames', case_sensitive=False, contains_search=True):
+    def search_cities(
+        self,
+        query: str,
+        attribute: CitySearchAttribute = 'alternatenames',
+        case_sensitive: bool = False,
+        contains_search: bool = True,
+    ) -> List[City]:
         """Search all city records and return list of records, that match query for given attribute."""
         results = []
         query = (case_sensitive and query) or query.casefold()
@@ -97,7 +111,7 @@ def search_cities(self, query, attribute='alternatenames', case_sensitive=False,
         return results
 
     @staticmethod
-    def _load_data(datadict, datafile):
+    def _load_data(datadict: Optional[Dict[Any, Any]], datafile: str) -> Dict[Any, Any]:
         if datadict is None:
             with open(os.path.join(os.path.dirname(__file__), 'data', datafile)) as f:
                 datadict = json.load(f)

geonamescache/geonamesdata.py

@@ -1,5 +1,9 @@
 # -*- coding: utf-8 -*-
-us_states = {
+from typing import Dict
+
+from .types import USState, USStateCode
+
+us_states: Dict[USStateCode, USState] = {
     'AK': {'code': 'AK', 'name': 'Alaska', 'fips': '02', 'geonameid': 5879092},
     'AL': {'code': 'AL', 'name': 'Alabama', 'fips': '01', 'geonameid': 4829764},
     'AR': {'code': 'AR', 'name': 'Arkansas', 'fips': '05', 'geonameid': 4099753},
@@ -51,4 +55,4 @@
     'WI': {'code': 'WI', 'name': 'Wisconsin', 'fips': '55', 'geonameid': 5279468},
     'WV': {'code': 'WV', 'name': 'West Virginia', 'fips': '54', 'geonameid': 4826850},
     'WY': {'code': 'WY', 'name': 'Wyoming', 'fips': '56', 'geonameid': 5843591}
-}
+}

geonamescache/mappers.py

@@ -1,9 +1,37 @@
 # -*- coding: utf-8 -*-
+from typing import Any, Callable, Literal, overload
+
 from geonamescache import GeonamesCache
+
 from . import mappings
+from .types import (ContinentCode, CountryFields, CountryNumericFields,
+                    CountryStringFields)
+
+
+@overload
+def country(
+    from_key: str = "name", *, to_key: CountryNumericFields
+) -> Callable[[str], int]: ...
+
+
+@overload
+def country(
+    from_key: str = "name",
+    to_key: CountryStringFields = "iso",
+) -> Callable[[str], str]: ...
+
+
+@overload
+def country(
+    from_key: str = "name",
+    *,
+    to_key: Literal["continentcode"],
+) -> Callable[[str], ContinentCode]: ...
 
 
-def country(from_key='name', to_key='iso'):
+def country(
+    from_key: str = "name", to_key: CountryFields = "iso"
+) -> Callable[[str], Any]:
     """Creates and returns a mapper function to access country data.
 
     The mapper function that is returned must be called with one argument. In
@@ -21,13 +49,13 @@ def country(from_key='name', to_key='iso'):
     gc = GeonamesCache()
     dataset = gc.get_dataset_by_key(gc.get_countries(), from_key)
 
-    def mapper(input):
+    def mapper(input: str) -> Any:
         # For country name inputs take the names mapping into account.
-        if 'name' == from_key:
+        if "name" == from_key:
             input = mappings.country_names.get(input, input)
         # If there is a record return the demanded attribute.
         item = dataset.get(input)
         if item:
             return item[to_key]
 
-    return mapper
+    return mapper

geonamescache/mappings.py

@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
+from typing import Dict
 
 # Map country name variants to the ones used in GeoNames.
-country_names = {
+country_names: Dict[str, str] = {
     'Bolivia (Plurinational State of)': 'Bolivia',
     'Bosnia-Herzegovina': 'Bosnia and Herzegovina',
     'Brunei Darussalam': 'Brunei',
@@ -72,4 +73,4 @@
     'Venezuela (Bolivarian Republic of)': 'Venezuela',
     'Viet Nam': 'Vietnam',
     'West Bank and Gaza Strip': 'Palestinian Territory'
-}
+}