diff --git a/geonamescache/__init__.py b/geonamescache/__init__.py index d43b23e..b87577b 100644 --- a/geonamescache/__init__.py +++ b/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) diff --git a/geonamescache/geonamesdata.py b/geonamescache/geonamesdata.py index a244815..271e2ce 100644 --- a/geonamescache/geonamesdata.py +++ b/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} -} \ No newline at end of file +} diff --git a/geonamescache/mappers.py b/geonamescache/mappers.py index 5879f1e..45d1aa4 100644 --- a/geonamescache/mappers.py +++ b/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 \ No newline at end of file + return mapper diff --git a/geonamescache/mappings.py b/geonamescache/mappings.py index 8d1c08a..d953a90 100644 --- a/geonamescache/mappings.py +++ b/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' -} \ No newline at end of file +} diff --git a/geonamescache/types.py b/geonamescache/types.py new file mode 100644 index 0000000..f1b033d --- /dev/null +++ b/geonamescache/types.py @@ -0,0 +1,238 @@ +# -*- 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 + + +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 + 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 diff --git a/requirements-dev.txt b/requirements-dev.txt index d506377..375ddc5 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,4 +5,5 @@ ipython ipdb pytest twine +typing-extensions wheel \ No newline at end of file diff --git a/setup.py b/setup.py index 077f356..431ba99 100644 --- a/setup.py +++ b/setup.py @@ -31,4 +31,5 @@ ], test_suite='tests', tests_require=['pytest'], + install_requires=["typing-extensions"] )