diff --git a/pycountry_wrapper/__init__.py b/pycountry_wrapper/__init__.py index ef89313..793fa56 100644 --- a/pycountry_wrapper/__init__.py +++ b/pycountry_wrapper/__init__.py @@ -1,6 +1,20 @@ from __future__ import annotations +from typing import Optional +from functools import wraps import pycountry + +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. @@ -9,45 +23,22 @@ class Country: - Country.from_alpha_2("DE") - Country.from_alpha_3("DEU") - If the country couldn't be found, it raises a ValueError. - """ - - def __new__(cls, country: str = None, pycountry_object = None, silent: bool = False): - self = super().__new__(cls) + 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: - if silent: - return None + + if pycountry_object is None and not allow_empty: raise ValueError(f"Country {country} couldn't be found") self.pycountry_object = pycountry_object - return self - - def __init__(self, country: str = None, pycountry_object = None, silent: bool = False) -> None: ... - - """ - def __init__(self, ): - 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 silent: - raise CountryDoesNotExist() - - self.pycountry_object = pycountry_object - - print("init") - """ @classmethod def from_alpha_2(cls, alpha_2: str) -> Country: @@ -61,6 +52,51 @@ class Country: 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 @@ -80,9 +116,3 @@ class Country: @property def official_name(self) -> 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__() diff --git a/pycountry_wrapper/__main__.py b/pycountry_wrapper/__main__.py index 78b0945..6ed4f72 100644 --- a/pycountry_wrapper/__main__.py +++ b/pycountry_wrapper/__main__.py @@ -5,4 +5,4 @@ from . import Country def cli(): print(f"Running {__name__} version {__version__} from __main__.py") - print(Country(country="DE")) \ No newline at end of file + print(Country(country="DE").name) \ No newline at end of file