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[str, TDict], key: str
+    ) -> Dict[str, 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[str, Any]], datafile: str) -> Dict[str, 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,12 @@
 # -*- coding: utf-8 -*-
+from typing import Callable, Union
+
 from geonamescache import GeonamesCache
+
 from . import mappings
 
 
-def country(from_key='name', to_key='iso'):
+def country(from_key: str = 'name', to_key: str = 'iso') -> Callable[[str], Union[str, int, None]]:
     """Creates and returns a mapper function to access country data.
 
     The mapper function that is returned must be called with one argument. In
@@ -21,7 +24,7 @@ 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) -> Union[str, int, None]:
         # For country name inputs take the names mapping into account.
         if 'name' == from_key:
             input = mappings.country_names.get(input, input)
@@ -30,4 +33,4 @@ def mapper(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'
-}
+}

geonamescache/types.py

@@ -0,0 +1,209 @@
+# -*- coding: utf-8 -*-
+from typing import List
+
+from typing_extensions import Literal, NotRequired, TypedDict
+
+GeoNameIdStr = str
+ISOStr = str
+ContinentCode = Literal['AF', 'AN', 'AS', 'EU', 'NA', 'OC', 'SA']
+USStateCode = Literal[
+    'AK',
+    'AL',
+    'AR',
+    'AZ',
+    'CA',
+    'CO',
+    'CT',
+    'DC',
+    'DE',
+    'FL',
+    'GA',
+    'HI',
+    'IA',
+    'ID',
+    'IL',
+    'IN',
+    'KS',
+    'KY',
+    'LA',
+    'MA',
+    'MD',
+    'ME',
+    'MI',
+    'MN',
+    'MO',
+    'MS',
+    'MT',
+    'NC',
+    'ND',
+    'NE',
+    'NH',
+    'NJ',
+    'NM',
+    'NV',
+    'NY',
+    'OH',
+    'OK',
+    'OR',
+    'PA',
+    'RI',
+    'SC',
+    'SD',
+    'TN',
+    'TX',
+    'UT',
+    'VA',
+    'VT',
+    'WA',
+    'WI',
+    'WV',
+    'WY',
+]
+USStateName = Literal[
+    'Alaska',
+    'Alabama',
+    'Arkansas',
+    'Arizona',
+    'California',
+    'Colorado',
+    'Connecticut',
+    'District of Columbia',
+    'Delaware',
+    'Florida',
+    'Georgia',
+    'Hawaii',
+    'Iowa',
+    'Idaho',
+    'Illinois',
+    'Indiana',
+    'Kansas',
+    'Kentucky',
+    'Louisiana',
+    'Massachusetts',
+    'Maryland',
+    'Maine',
+    'Michigan',
+    'Minnesota',
+    'Missouri',
+    'Mississippi',
+    'Montana',
+    'North Carolina',
+    'North Dakota',
+    'Nebraska',
+    'New Hampshire',
+    'New Jersey',
+    'New Mexico',
+    'Nevada',
+    'New York',
+    'Ohio',
+    'Oklahoma',
+    'Oregon',
+    'Pennsylvania',
+    'Rhode Island',
+    'South Carolina',
+    'South Dakota',
+    'Tennessee',
+    'Texas',
+    'Utah',
+    'Virginia',
+    'Vermont',
+    'Washington',
+    'Wisconsin',
+    'West Virginia',
+    'Wyoming',
+]
+CitySearchAttribute = Literal['alternatenames', 'admin1code', 'countrycode', 'name', 'timezone']
+
+
+class TimeZone(TypedDict):
+    dstOffset: int
+    gmtOffset: int
+    timeZoneId: str
+
+
+class BBox(TypedDict):
+    accuracyLevel: int
+    east: float
+    north: float
+    south: float
+    west: float
+
+
+class ContinentAlternateName(TypedDict):
+    lang: str
+    name: str
+    isPreferredName: NotRequired[bool]
+    isShortName: NotRequired[bool]
+    isColloquial: NotRequired[bool]
+
+
+class Continent(TypedDict):
+    alternateNames: List[ContinentAlternateName]
+    adminName1: str
+    adminName2: str
+    adminName3: str
+    adminName4: str
+    adminName5: str
+    asciiName: str
+    continentCode: ContinentCode
+    fclName: str
+    fcodeName: str
+    astergdem: int
+    bbox: BBox
+    geonameId: int
+    fcl: str
+    fcode: str
+    lat: str
+    lng: str
+    name: str
+    population: int
+    timezone: TimeZone
+    toponymName: str
+    srtm3: int
+    wikipediaURL: str
+    cc2: NotRequired[str]
+
+
+class City(TypedDict):
+    alternatenames: List[str]
+    admin1code: str
+    countrycode: str
+    geonameid: int
+    latitude: float
+    longitude: float
+    name: str
+    population: int
+    timezone: str
+
+
+class Country(TypedDict):
+    areakm2: int
+    capital: str
+    continentcode: ContinentCode
+    currencycode: str
+    currencyname: str
+    iso: str
+    isonumeric: int
+    iso3: str
+    fips: str
+    geonameid: int
+    languages: str
+    name: str
+    neighbours: str
+    phone: str
+    population: int
+    postalcoderegex: str
+    tld: str
+
+
+class USState(TypedDict):
+    code: USStateCode
+    fips: str
+    geonameid: int
+    name: USStateName
+
+
+class USCounty(TypedDict):
+    fips: str
+    name: str
+    state: USStateCode

requirements-dev.txt

@@ -5,4 +5,5 @@ ipython
 ipdb
 pytest
 twine
+typing-extensions
 wheel

setup.py

@@ -31,4 +31,5 @@
     ],
     test_suite='tests',
     tests_require=['pytest'],
+    install_requires=["typing-extensions"]
 )