diff --git a/geonamescache/__init__.py b/geonamescache/__init__.py index a781ed2..b87577b 100644 --- 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 index 5042793..45d1aa4 100644 --- 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 index 38dea39..f1b033d 100644 --- 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