Merge branch 'typehints'

yaph · Jul 25, 2024 · e17902b · e17902b
2 parents e366187 + ff92eb4
commit e17902b
Show file tree

Hide file tree

Showing 3 changed files with 165 additions and 111 deletions.
diff --git a/geonamescache/__init__.py b/geonamescache/__init__.py
@@ -31,8 +31,8 @@ def __init__(self, min_city_population: int = 15000):
         self.min_city_population = min_city_population
 
     def get_dataset_by_key(
-        self, dataset: Dict[str, TDict], key: str
-    ) -> Dict[str, TDict]:
+        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) -> Dict[ContinentCode, Continent]:
@@ -111,7 +111,7 @@ def search_cities(
         return results
 
     @staticmethod
-    def _load_data(datadict: Optional[Dict[str, Any]], datafile: str) -> Dict[str, Any]:
+    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)

diff --git a/geonamescache/mappers.py b/geonamescache/mappers.py
@@ -1,12 +1,37 @@
 # -*- coding: utf-8 -*-
-from typing import Callable, Union
+from typing import Any, Callable, Literal, overload
 
 from geonamescache import GeonamesCache
 
 from . import mappings
+from .types import (ContinentCode, CountryFields, CountryNumericFields,
+                    CountryStringFields)
 
 
-def country(from_key: str = 'name', to_key: str = 'iso') -> Callable[[str], Union[str, int, None]]:
+@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: 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
@@ -24,9 +49,9 @@ def country(from_key: str = 'name', to_key: str = 'iso') -> Callable[[str], Unio
     gc = GeonamesCache()
     dataset = gc.get_dataset_by_key(gc.get_countries(), from_key)
 
-    def mapper(input: str) -> Union[str, int, None]:
+    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)

diff --git a/geonamescache/types.py b/geonamescache/types.py
@@ -5,114 +5,116 @@
 
 GeoNameIdStr = str
 ISOStr = str
-ContinentCode = Literal['AF', 'AN', 'AS', 'EU', 'NA', 'OC', 'SA']
+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',
+    "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',
+    "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"
 ]
-CitySearchAttribute = Literal['alternatenames', 'admin1code', 'countrycode', 'name', 'timezone']
 
 
 class TimeZone(TypedDict):
@@ -176,6 +178,33 @@ class City(TypedDict):
     timezone: str
 
 
+CountryNumericFields = Literal[
+    "areakm2",
+    "isonumeric",
+    "geonameid",
+    "population",
+]
+CountryStringFields = Literal[
+    "capital",
+    "currencycode",
+    "currencyname",
+    "iso",
+    "iso3",
+    "fips",
+    "languages",
+    "name",
+    "neighbours",
+    "phone",
+    "postalcoderegex",
+    "tld",
+]
+CountryFields = Literal[
+    CountryNumericFields,
+    CountryStringFields,
+    "continentcode",
+]
+
+
 class Country(TypedDict):
     areakm2: int
     capital: str