Remaining type hints improvements

Author: spjuhel

climada/entity/impact_funcs/base.py

@@ -189,7 +189,7 @@ def from_step_impf(
         haz_type: str,
         mdd: tuple[float, float] = (0, 1),
         paa: tuple[float, float] = (1, 1),
-        impf_id: int = 1,
+        impf_id: int | str = 1,
         **kwargs,
     ):
         """Step function type impact function.
@@ -207,7 +207,7 @@ def from_step_impf(
             (min, max) mdd values. The default is (0, 1)
         paa: tuple(float, float)
             (min, max) paa values. The default is (1, 1)
-        impf_id : int, optional, default=1
+        impf_id : int|str, optional, default=1
             impact function id
         kwargs :
             keyword arguments passed to ImpactFunc()
@@ -250,7 +250,7 @@ def from_sigmoid_impf(
         k: float,
         x0: float,
         haz_type: str,
-        impf_id: int = 1,
+        impf_id: int | str = 1,
         **kwargs,
     ):
         r"""Sigmoid type impact function hinging on three parameter.
@@ -320,7 +320,7 @@ def from_poly_s_shape(
         scale: float,
         exponent: float,
         haz_type: str,
-        impf_id: int = 1,
+        impf_id: int | str = 1,
         **kwargs,
     ):
         r"""S-shape polynomial impact function hinging on four parameter.

climada/entity/impact_funcs/impact_func_set.py

@@ -141,7 +141,9 @@ def append(self, func: ImpactFunc):
             self._data[func.haz_type] = dict()
         self._data[func.haz_type][func.id] = func
 
-    def remove_func(self, haz_type: str = None, fun_id: str | int = None):
+    def remove_func(
+        self, haz_type: Optional[str] = None, fun_id: Optional[str | int] = None
+    ):
         """Remove impact function(s) with provided hazard type and/or id.
         If no input provided, all impact functions are removed.
 
@@ -174,21 +176,21 @@ def remove_func(self, haz_type: str = None, fun_id: str | int = None):
             self._data = dict()
 
     @overload
-    def get_func(self, haz_type: str, fun_id: int) -> ImpactFunc: ...
+    def get_func(self, haz_type: str, fun_id: int | str) -> ImpactFunc: ...
 
     @overload
     def get_func(self, haz_type: str, fun_id: None = None) -> list[ImpactFunc]: ...
 
     @overload
-    def get_func(self, haz_type: None, fun_id: int) -> list[ImpactFunc]: ...
+    def get_func(self, haz_type: None, fun_id: int | str) -> list[ImpactFunc]: ...
 
     @overload
     def get_func(
         self, haz_type: None = None, fun_id: None = None
     ) -> Dict[str, Dict[int, ImpactFunc]]: ...
 
     def get_func(
-        self, haz_type: Optional[str] = None, fun_id: Optional[int] = None
+        self, haz_type: Optional[str] = None, fun_id: Optional[int | str] = None
     ) -> Union[ImpactFunc, list[ImpactFunc], Dict[str, Dict[int, ImpactFunc]]]:
         """Get ImpactFunc(s) of input hazard type and/or id.
         If no input provided, all impact functions are returned.
@@ -225,7 +227,7 @@ def get_func(
         else:
             return self._data
 
-    def get_hazard_types(self, fun_id: str | int) -> list[str]:
+    def get_hazard_types(self, fun_id: Optional[str | int]) -> list[str]:
         """Get impact functions hazard types contained for the id provided.
         Return all hazard types if no input id.
 
@@ -247,7 +249,15 @@ def get_hazard_types(self, fun_id: str | int) -> list[str]:
                 haz_types.append(vul_haz)
         return haz_types
 
-    def get_ids(self, haz_type: str = None) -> list[int | str]:
+    @overload
+    def get_ids(self, haz_type: None = None) -> dict[str, list[str | int]]: ...
+
+    @overload
+    def get_ids(self, haz_type: str) -> list[int | str]: ...
+
+    def get_ids(
+        self, haz_type: Optional[str] = None
+    ) -> dict[str, list[str | int]] | list[int | str]:
         """Get impact functions ids contained for the hazard type provided.
         Return all ids for each hazard type if no input hazard type.
 
@@ -272,7 +282,9 @@ def get_ids(self, haz_type: str = None) -> list[int | str]:
         except KeyError:
             return list()
 
-    def size(self, haz_type: str = None, fun_id: str | int = None) -> int:
+    def size(
+        self, haz_type: Optional[str] = None, fun_id: Optional[str | int] = None
+    ) -> int:
         """Get number of impact functions contained with input hazard type and
         /or id. If no input provided, get total number of impact functions.
 
@@ -295,6 +307,7 @@ def size(self, haz_type: str = None, fun_id: str | int = None) -> int:
             return 1
         if (haz_type is not None) or (fun_id is not None):
             return len(self.get_func(haz_type, fun_id))
+
         return sum(len(vul_list) for vul_list in self.get_ids().values())
 
     def check(self):
@@ -316,7 +329,7 @@ def check(self):
                     )
                 vul.check()
 
-    def extend(self, impact_funcs: ImpactFuncSet):
+    def extend(self, impact_funcs: "ImpactFuncSet"):
         """Append impact functions of input ImpactFuncSet to current
         ImpactFuncSet. Overwrite ImpactFunc if same id and haz_type.
 
@@ -339,7 +352,13 @@ def extend(self, impact_funcs: ImpactFuncSet):
             for _, vul in vul_dict.items():
                 self.append(vul)
 
-    def plot(self, haz_type: str = None, fun_id: str | int = None, axis=None, **kwargs):
+    def plot(
+        self,
+        haz_type: Optional[str] = None,
+        fun_id: Optional[str | int] = None,
+        axis=None,
+        **kwargs,
+    ):
         """Plot impact functions of selected hazard (all if not provided) and
         selected function id (all if not provided).