explicitly type country field types

yaph · Jul 25, 2024 · 89ae766 · 89ae766
1 parent d0cc123
commit 89ae766
Show file tree

Hide file tree

Showing 2 changed files with 139 additions and 115 deletions.
diff --git a/geonamescache/mappers.py b/geonamescache/mappers.py
@@ -4,19 +4,12 @@
 from geonamescache import GeonamesCache
 
 from . import mappings
-from .types import ContinentCode
+from .types import ContinentCode, CountryFields, CountryNumericFields
 
 
 @overload
 def country(
-    from_key: str = "name",
-    *,
-    to_key: Literal[
-        "areakm2",
-        "isonumeric",
-        "geonameid",
-        "population",
-    ],
+    from_key: str = "name", *, to_key: CountryNumericFields
 ) -> Callable[[str], int]: ...
 
 
@@ -48,7 +41,9 @@ def country(
 ) -> Callable[[str], ContinentCode]: ...
 
 
-def country(from_key: str = "name", to_key: str = "iso") -> Callable[[str], Any]:
+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
@@ -73,6 +68,6 @@ def mapper(input: str) -> Any:
         # If there is a record return the demanded attribute.
         item = dataset.get(input)
         if item:
-            return item[to_key]  # type: ignore
+            return item[to_key]
 
     return mapper
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