output type annotations in code utils

src/hydamo/hydamo/code_utils.py

@@ -13,7 +13,7 @@
 NEN3610_ID_TEMPLATE = "NL.BGTCODE.{bgt_code}.{layer}.{nen3610_id}"
 
 
-def get_codes_df():
+def get_codes_df() -> DataFrame:
     """Get (and set) the global CODES_DF.
 
     Returns
@@ -33,19 +33,19 @@ def get_codes_df():
     return CODES_DF.copy()
 
 
-def bgt_code_exists(bgt_code: str):
+def bgt_code_exists(bgt_code: str) -> bool:
     """Check if bgt_code exists in CODES_DF"""
     codes_df = get_codes_df()
     return bgt_code in codes_df.bgt_code.to_numpy()
 
 
-def wbh_code_exists(wbh_code):
+def wbh_code_exists(wbh_code) -> bool:
     """Check if wbh_code exists in CODES_DF"""
     codes_df = get_codes_df()
     return wbh_code in codes_df.wbh_code.to_numpy()
 
 
-def bgt_to_wbh_code(bgt_code):
+def bgt_to_wbh_code(bgt_code) -> Union[str, None]:
     """Convert bgt_code to wbh_code if bgt_code exists"""
     if bgt_code_exists(bgt_code):
         codes_df = get_codes_df()
@@ -58,7 +58,7 @@ def find_codes(
     organization: str,
     administration_category: Union[str, None] = None,
     to_dict: bool = True,
-):
+) -> dict:
     """Find codes associated with an organization"""
     codes_df = get_codes_df()
 
@@ -93,7 +93,7 @@ def find_codes(
             return df.to_dict()
 
 
-def generate_model_id(code, layer, wbh_code=None, bgt_code=None, geometry=None):
+def generate_model_id(code, layer, wbh_code=None, bgt_code=None, geometry=None) -> str:
     """Generate a model_id from wbh_code or bgt_code and code or x/y coordinate"""
     if code is None:
         if geometry is not None:
@@ -111,14 +111,18 @@ def generate_model_id(code, layer, wbh_code=None, bgt_code=None, geometry=None):
     result = None
     if wbh_code:
         if wbh_code_exists(wbh_code):
-            result = WBH_CODE_TEMPLATE.format(wbh_code=wbh_code, code=code)
+            result = WBH_CODE_TEMPLATE.format(wbh_code=wbh_code, layer=layer, code=code)
     elif bgt_code:
         if bgt_code_exists(bgt_code):
             wbh_code = bgt_to_wbh_code(bgt_code)
             if wbh_code:
-                result = WBH_CODE_TEMPLATE.format(wbh_code=wbh_code, code=code)
+                result = WBH_CODE_TEMPLATE.format(
+                    wbh_code=wbh_code, layer=layer, code=code
+                )
             else:
-                result = BGT_CODE_TEMPLATE.format(bgt_code=bgt_code, code=code)
+                result = BGT_CODE_TEMPLATE.format(
+                    bgt_code=bgt_code, layer=layer, code=code
+                )
     if result is None:
         raise ValueError(
             f"""

src/hydamo/tests/test_code_utils.py

@@ -20,9 +20,12 @@ def test_generate_model_id():
     # see if we can find Zuid-Holland
     result = code_utils.find_codes("Zuid-Holland")
     assert result["wbh_code"] == code_utils.bgt_to_wbh_code(result["bgt_code"])
-
-    model_id = code_utils.generate_model_id(code="test", bgt_code=result["bgt_code"])
-    assert model_id == "NL.WBHCODE.69.test"
+    code = "KGM001"
+    layer = "gemaal"
+    model_id = code_utils.generate_model_id(
+        code=code, layer=layer, bgt_code=result["bgt_code"]
+    )
+    assert model_id == f"NL.WBHCODE.69.{layer}.{code}"
     assert model_id == code_utils.generate_model_id(
-        code="test", wbh_code=result["wbh_code"]
+        code=code, layer=layer, wbh_code=result["wbh_code"]
     )