feat: added better typing

This commit is contained in:
Hazel 2024-11-29 16:52:35 +01:00
parent 4964d5a494
commit bac32a116e
2 changed files with 66 additions and 36 deletions

View File

@ -1,6 +1,20 @@
from __future__ import annotations from __future__ import annotations
from typing import Optional
from functools import wraps
import pycountry 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: class Country:
""" """
This gets countries based on the ISO 3166-1 standart. This gets countries based on the ISO 3166-1 standart.
@ -9,45 +23,22 @@ class Country:
- Country.from_alpha_2("DE") - Country.from_alpha_2("DE")
- Country.from_alpha_3("DEU") - Country.from_alpha_3("DEU")
If the country couldn't be found, it raises a ValueError. If the country couldn't be found, it raises a ValueError, or creates an empty object.
""" Empty objects return for every attribute None
"""
def __new__(cls, country: str = None, pycountry_object = None, silent: bool = False):
self = super().__new__(cls)
def __init__(self, country: Optional[str] = None, pycountry_object = None, allow_empty: bool = True) -> None:
if country is not None: if country is not None:
# auto detect if alpha_2 or alpha_3 # auto detect if alpha_2 or alpha_3
if len(country) == 2: if len(country) == 2:
pycountry_object = pycountry.countries.get(alpha_2=country.upper()) pycountry_object = pycountry.countries.get(alpha_2=country.upper())
elif len(country) == 3: elif len(country) == 3:
pycountry_object = pycountry.countries.get(alpha_3=country.upper()) pycountry_object = pycountry.countries.get(alpha_3=country.upper())
if pycountry_object is None: if pycountry_object is None and not allow_empty:
if silent:
return None
raise ValueError(f"Country {country} couldn't be found") raise ValueError(f"Country {country} couldn't be found")
self.pycountry_object = pycountry_object 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 @classmethod
def from_alpha_2(cls, alpha_2: str) -> Country: def from_alpha_2(cls, alpha_2: str) -> Country:
@ -61,6 +52,51 @@ class Country:
def from_fuzzy(cls, fuzzy: str) -> Country: def from_fuzzy(cls, fuzzy: str) -> Country:
return cls(pycountry_object=pycountry.countries.search_fuzzy(fuzzy)) 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 @property
def name(self) -> str: def name(self) -> str:
return self.pycountry_object.name return self.pycountry_object.name
@ -80,9 +116,3 @@ class Country:
@property @property
def official_name(self) -> str: def official_name(self) -> str:
return self.pycountry_object.official_name return self.pycountry_object.official_name
def __str__(self) -> str:
return self.pycountry_object.__str__()
def __repr__(self) -> str:
return self.pycountry_object.__repr__()

View File

@ -5,4 +5,4 @@ from . import Country
def cli(): def cli():
print(f"Running {__name__} version {__version__} from __main__.py") print(f"Running {__name__} version {__version__} from __main__.py")
print(Country(country="DE")) print(Country(country="DE").name)