From bac32a116ea9c0f7f76fe69aebc4c0f33061da90 Mon Sep 17 00:00:00 2001
From: Lars Noack <lars@webcontact.de>
Date: Fri, 29 Nov 2024 16:52:35 +0100
Subject: [PATCH] feat: added better typing

---
 pycountry_wrapper/__init__.py | 100 ++++++++++++++++++++++------------
 pycountry_wrapper/__main__.py |   2 +-
 2 files changed, 66 insertions(+), 36 deletions(-)

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