Suggested fix for ImpfSetTropCyclone.get_countries_per_region() #1105

climada/entity/impact_funcs/test/test_tc.py

@@ -24,7 +24,11 @@
 import numpy as np
 import pandas as pd
 
-from climada.entity.impact_funcs.trop_cyclone import ImpfSetTropCyclone, ImpfTropCyclone
+from climada.entity.impact_funcs.trop_cyclone import (
+    CountryCode,
+    ImpfSetTropCyclone,
+    ImpfTropCyclone,
+)
 
 
 class TestEmanuelFormula(unittest.TestCase):
@@ -168,6 +172,16 @@ def test_get_countries_per_region(self):
         self.assertListEqual(out[2], [124, 840])
         self.assertListEqual(out[3], ["CAN", "USA"])
 
+    def test_get_countries_per_region_all_or_none(self):
+        ifs = ImpfSetTropCyclone()
+        out = ifs.get_countries_per_region()
+        out2 = ifs.get_countries_per_region("all")
+        self.assertEqual(out, out2)
+        for reg in CountryCode.REGION_NAME.value:
+            out_reg = ifs.get_countries_per_region(reg)
+            for i in range(4):
+                self.assertEqual(out[i][reg], out_reg[i])
+
     def test_get_imf_id_regions_per_countries(self):
         """Test get_impf_id_regions_per_countries()"""
         ifs = ImpfSetTropCyclone()

climada/entity/impact_funcs/trop_cyclone.py

@@ -23,6 +23,7 @@
 
 import logging
 from enum import Enum
+from typing import Literal, overload
 
 import numpy as np
 import pandas as pd
@@ -92,7 +93,7 @@ class CountryCode(Enum):
             "NAM", "NER", "NGA", "NLD", "NOR", "POL", "PRK", "PRT", "PSE", "REU", "ROU",
             "RUS", "RWA", "SDN", "SEN", "SGP", "SGS", "SJM", "SLE", "SMR", "SPM", "SRB",
             "SSD", "STP", "SVK", "SVN", "SWE", "SYC", "TCD", "TGO", "TUN", "TUR", "UKR",
-            "UMI", "VAT", "XKX", "ZMB",
+            "UMI", "VAT", "XKO", "ZMB",
         ],
     }
 
@@ -363,10 +364,47 @@ def calibrated_regional_vhalf(
             reg_v_half[regions_short[-1]] = np.round(df_reg["v_half"].values[0], 5)
         return reg_v_half
 
+    @overload
+    @staticmethod
+    def get_countries_per_region(
+        region: Literal["all"] = "all",
+    ) -> tuple[
+        dict[str, str],  # region_name
+        dict[str, int],  # impf_id
+        dict[str, list[int]],  # numeric
+        dict[str, list[str]],  # alpha3
+    ]: ...
+
+    @overload
+    @staticmethod
+    def get_countries_per_region(
+        region: None,
+    ) -> tuple[
+        dict[str, str], dict[str, int], dict[str, list[int]], dict[str, list[str]]
+    ]: ...
+
+    @overload
+    @staticmethod
+    def get_countries_per_region(
+        region: str,
+    ) -> tuple[
+        str, int, list[int], list[str]  # region_name  # impf_id  # numeric  # alpha3
+    ]: ...
+
     @staticmethod
     def get_countries_per_region(region=None):
-        """Returns dictionaries with numerical (numeric) and alphabetical (alpha3) ISO3 codes
-        of all countries associated to a calibration region.
+        """Returns countries within a TC calibration region and associated impact functions.
+
+        This method returns a tuple with numerical (numeric) and alphabetical (alpha3)
+        ISO3 codes of all countries associated to a calibration region.
+
+        If no region or "all" is provided as argument, the method return a tuple of
+        dictionaries with short name of the tropical cyclone calibration regions as
+        keys and the values for each of those.
+
+        Notes
+        -----
+
         Only contains countries that were affected by tropical cyclones
         between 1980 and 2017 according to EM-DAT.
 
@@ -395,9 +433,12 @@ def get_countries_per_region(region=None):
             return (
                 CountryCode.REGION_NAME.value,
                 CountryCode.IMPF_ID.value,
-                coordinates.country_to_iso(
-                    CountryCode.ALPHA3.value, representation="numeric"
-                ),
+                {
+                    reg: coordinates.country_to_iso(
+                        CountryCode.ALPHA3.value[reg], representation="numeric"
+                    )
+                    for reg in CountryCode.REGION_NAME.value
+                },
                 CountryCode.ALPHA3.value,
             )