diff --git a/pycountry_wrapper/__init__.py b/pycountry_wrapper/__init__.py index 793fa56..257425a 100644 --- a/pycountry_wrapper/__init__.py +++ b/pycountry_wrapper/__init__.py @@ -1,118 +1,8 @@ -from __future__ import annotations -from typing import Optional -from functools import wraps -import pycountry +from .country import Country, StrictCountry +from . import config - -def none_if_empty(func): - @wraps(func) - def wrapper(self, *args, **kwargs): - if self.is_empty: - return None - - return func(self, *args, **kwargs) - - return wrapper - - -class Country: - """ - This gets countries based on the ISO 3166-1 standart. - - Two examples are: - - Country.from_alpha_2("DE") - - Country.from_alpha_3("DEU") - - If the country couldn't be found, it raises a ValueError, or creates an empty object. - Empty objects return for every attribute None - """ - - def __init__(self, country: Optional[str] = None, pycountry_object = None, allow_empty: bool = True) -> None: - if country is not None: - # auto detect if alpha_2 or alpha_3 - if len(country) == 2: - pycountry_object = pycountry.countries.get(alpha_2=country.upper()) - elif len(country) == 3: - pycountry_object = pycountry.countries.get(alpha_3=country.upper()) - - if pycountry_object is None and not allow_empty: - raise ValueError(f"Country {country} couldn't be found") - - self.pycountry_object = pycountry_object - - @classmethod - def from_alpha_2(cls, alpha_2: str) -> Country: - return cls(pycountry_object=pycountry.countries.get(alpha_2=alpha_2.upper())) - - @classmethod - def from_alpha_3(cls, alpha_3: str) -> Country: - return cls(pycountry_object=pycountry.countries.get(alpha_3=alpha_3.upper())) - - @classmethod - def from_fuzzy(cls, fuzzy: str) -> Country: - return cls(pycountry_object=pycountry.countries.search_fuzzy(fuzzy)) - - @property - def is_empty(self) -> bool: - return self.pycountry_object is None - - @property - @none_if_empty - def name(self) -> Optional[str]: - return self.pycountry_object.name - - @property - @none_if_empty - def alpha_2(self) -> Optional[str]: - return self.pycountry_object.alpha_2 - - @property - @none_if_empty - def alpha_3(self) -> Optional[str]: - return self.pycountry_object.alpha_3 - - @property - @none_if_empty - def numeric(self) -> Optional[str]: - return self.pycountry_object.numeric - - @property - @none_if_empty - def official_name(self) -> Optional[str]: - return self.pycountry_object.official_name - - def __str__(self) -> str: - return self.pycountry_object.__str__() - - def __repr__(self) -> str: - return self.pycountry_object.__repr__() - - -class StrictCountry(Country): - """ - This works just like Country, - but the object cant be empty - """ - - def __init__(self, country: Optional[str] = None, pycountry_object = None) -> None: - super().__init__(country=country, pycountry_object=pycountry_object, allow_empty=False) - - @property - def name(self) -> str: - return self.pycountry_object.name - - @property - def alpha_2(self) -> str: - return self.pycountry_object.alpha_2 - - @property - def alpha_3(self) -> str: - return self.pycountry_object.alpha_3 - - @property - def numeric(self) -> str: - return self.pycountry_object.numeric - - @property - def official_name(self) -> str: - return self.pycountry_object.official_name +__all__ = [ + "Country", + "StrictCountry", + "config", +] diff --git a/pycountry_wrapper/config.py b/pycountry_wrapper/config.py new file mode 100644 index 0000000..91ced91 --- /dev/null +++ b/pycountry_wrapper/config.py @@ -0,0 +1,8 @@ +import typing as _t + +# defines the fallback country if a country can't be found +# alpha_2 or alpha_3 of ISO 3166-1 +fallback_country: _t.Optional[str] = None + +# should use fuzzy search if it cant find the country with alpha_2 or alpha_3 +allow_fuzzy_search: bool = True diff --git a/pycountry_wrapper/country.py b/pycountry_wrapper/country.py new file mode 100644 index 0000000..845ddd7 --- /dev/null +++ b/pycountry_wrapper/country.py @@ -0,0 +1,151 @@ +from __future__ import annotations +from typing import Optional +from functools import wraps +import pycountry + +from . import config + + +def none_if_empty(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + if self.is_empty: + return None + + return func(self, *args, **kwargs) + + return wrapper + + +class Country: + """ + This gets countries based on the ISO 3166-1 standart. + + Two examples are: + - Country.from_alpha_2("DE") + - Country.from_alpha_3("DEU") + + If the country couldn't be found, it raises a ValueError, or creates an empty object. + Empty objects return for every attribute None + """ + + def __init__(self, country: Optional[str] = None, pycountry_object = None) -> None: + self.pycountry_object = pycountry_object + + if self.pycountry_object is None: + # search for the country string instead if the pycountry_object isn't given + # this also implements the optional fallback + self.pycountry_object = self._search_pycountry_object(country=country) + + + if pycountry_object is None: + raise ValueError(f"Country {country} couldn't be found") + + + @classmethod + def _search_pycountry_object(cls, country: Optional[str], is_fallback: bool = False): + # fallback to configured country if necessary + if country is None: + if is_fallback: + return None + + return cls._search_pycountry_object(country=config.fallback_country, is_fallback=True) + + pycountry_object = None + + # the reason I don't immediately return the result is because then there would be a chance + # I would return None even though a country could be found through fuzzy search + country = country.strip() + if len(country) == 2: + pycountry_object = pycountry.countries.get(alpha_2=country.upper()) + elif len(country) == 3: + pycountry_object = pycountry.countries.get(alpha_3=country.upper()) + if pycountry_object is not None: + return cls(pycountry_object=pycountry_object) + + # fuzzy search if enabled + if config.allow_fuzzy_search: + found_countries = pycountry.countries.search_fuzzy(country) + if len(found_countries): + return cls(pycountry_object=found_countries[0]) + + @classmethod + def search(cls, country: Optional[str]) -> Optional[Country]: + return cls(pycountry_object=cls._search_pycountry_object(country=country)) + + @classmethod + def from_alpha_2(cls, alpha_2: str) -> Country: + return cls(pycountry_object=pycountry.countries.get(alpha_2=alpha_2.upper())) + + @classmethod + def from_alpha_3(cls, alpha_3: str) -> Country: + return cls(pycountry_object=pycountry.countries.get(alpha_3=alpha_3.upper())) + + @classmethod + def from_fuzzy(cls, fuzzy: str) -> Country: + return cls(pycountry_object=pycountry.countries.search_fuzzy(fuzzy)) + + @property + def is_empty(self) -> bool: + return self.pycountry_object is None + + @property + @none_if_empty + def name(self) -> Optional[str]: + return self.pycountry_object.name + + @property + @none_if_empty + def alpha_2(self) -> Optional[str]: + return self.pycountry_object.alpha_2 + + @property + @none_if_empty + def alpha_3(self) -> Optional[str]: + return self.pycountry_object.alpha_3 + + @property + @none_if_empty + def numeric(self) -> Optional[str]: + return self.pycountry_object.numeric + + @property + @none_if_empty + def official_name(self) -> Optional[str]: + return self.pycountry_object.official_name + + def __str__(self) -> str: + return self.pycountry_object.__str__() + + def __repr__(self) -> str: + return self.pycountry_object.__repr__() + + +class StrictCountry(Country): + """ + This works just like Country, + but the object cant be empty + """ + + def __init__(self, country: Optional[str] = None, pycountry_object = None) -> None: + super().__init__(country=country, pycountry_object=pycountry_object, allow_empty=False) + + @property + def name(self) -> str: + return self.pycountry_object.name + + @property + def alpha_2(self) -> str: + return self.pycountry_object.alpha_2 + + @property + def alpha_3(self) -> str: + return self.pycountry_object.alpha_3 + + @property + def numeric(self) -> str: + return self.pycountry_object.numeric + + @property + def official_name(self) -> str: + return self.pycountry_object.official_name