pycountry-wrapper/pycountry_wrapper/__init__.py

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.

    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