Fix behavior: Make time unit required in time conversion, not in regex.

Author: MImmesberger

src/ttsim/compute_taxes_and_transfers.py

@@ -196,14 +196,17 @@ def _get_top_level_namespace(
         if match := re_pattern.fullmatch(element):
             function_base_name = match.group("base_name")
             create_conversions_for_time_units = bool(match.group("time_unit"))
-
-            potential_derived_functions = potential_target_names_from_base_name(
-                base_name=function_base_name,
-                supported_time_conversions=supported_time_conversions,
-                supported_groupings=supported_groupings,
-                create_conversions_for_time_units=create_conversions_for_time_units,
-            )
-            potential_function_names.update(potential_derived_functions)
+        else:
+            function_base_name = element
+            create_conversions_for_time_units = False
+
+        potential_derived_functions = potential_target_names_from_base_name(
+            base_name=function_base_name,
+            supported_time_conversions=supported_time_conversions,
+            supported_groupings=supported_groupings,
+            create_conversions_for_time_units=create_conversions_for_time_units,
+        )
+        potential_function_names.update(potential_derived_functions)
 
     return potential_function_names
 

src/ttsim/shared.py

@@ -42,10 +42,11 @@ def get_re_pattern_for_time_units_and_groupings(
 ) -> re.Pattern:
     """Get a regex pattern for time units and groupings.
 
-    The pattern is of the form:
-    <base_name>_<time_unit>_<aggregation>
-    where <base_name> is optional, <time_unit> is one of the supported time units, and
-    <aggregation> is one of the supported groupings.
+    The pattern matches strings in any of these formats:
+    - <base_name>  (can contain underscores)
+    - <base_name>_<time_unit>
+    - <base_name>_<aggregation>
+    - <base_name>_<time_unit>_<aggregation>
 
     Parameters
     ----------
@@ -60,9 +61,12 @@ def get_re_pattern_for_time_units_and_groupings(
         The regex pattern.
     """
     units = "".join(supported_time_units)
-    groupings = "|".join([f"_{grouping}" for grouping in supported_groupings])
+    groupings = "|".join(supported_groupings)
     return re.compile(
-        f"(?P<base_name>.*_)(?P<time_unit>[{units}])(?P<aggregation>{groupings})?"
+        f"(?P<base_name>.*?)"
+        f"(?:_(?P<time_unit>[{units}]))?"
+        f"(?:_(?P<aggregation>{groupings}))?"
+        f"$"
     )
 
 

src/ttsim/time_conversion.py

@@ -453,17 +453,22 @@ def _create_time_conversion_functions(
         supported_groupings=SUPPORTED_GROUPINGS,
         supported_time_units=all_time_units,
     )
+
     match = time_unit_pattern.fullmatch(name)
-    dependencies = set(inspect.signature(func).parameters) if func else set()
+    base_name = match.group("base_name") or ""
+    time_unit = match.group("time_unit") or ""
+    aggregation = match.group("aggregation") or ""
 
-    if match:
-        base_name = match.group("base_name")
-        time_unit = match.group("time_unit")
-        aggregation = match.group("aggregation") or ""
+    dependencies = set(inspect.signature(func).parameters) if func else set()
 
+    if match and time_unit:
         missing_time_units = [unit for unit in all_time_units if unit != time_unit]
         for missing_time_unit in missing_time_units:
-            new_name = f"{base_name}{missing_time_unit}{aggregation}"
+            new_name = (
+                f"{base_name}_{missing_time_unit}{aggregation}"
+                if aggregation
+                else f"{base_name}_{missing_time_unit}"
+            )
 
             # Without this check, we could create cycles in the DAG: Consider a
             # hard-coded function `var_y` that takes `var_m` as an input, assuming it

tests/ttsim/test_compute_taxes_and_transfers.py

@@ -16,6 +16,7 @@
     _fail_if_foreign_keys_are_invalid,
     _fail_if_group_variables_not_constant_within_groups,
     _fail_if_pid_is_non_unique,
+    _get_top_level_namespace,
     _partial_parameters_to_functions,
     compute_taxes_and_transfers,
 )
@@ -699,3 +700,58 @@ def test_fail_if_cannot_be_converted_to_correct_type(
 def test_assert_valid_ttsim_pytree(tree, leaf_checker, err_substr):
     with pytest.raises(TypeError, match=re.escape(err_substr)):
         assert_valid_ttsim_pytree(tree, leaf_checker, "tree")
+
+
+@pytest.mark.parametrize(
+    (
+        "environment",
+        "supported_time_conversions",
+        "supported_groupings",
+        "expected",
+    ),
+    [
+        (
+            PolicyEnvironment(
+                raw_objects_tree={
+                    "foo_m": policy_function(leaf_name="foo_m")(lambda x: x)
+                },
+                aggregation_specs_tree={},
+            ),
+            ["m", "y"],
+            ["hh"],
+            {"foo_m", "foo_y", "foo_m_hh", "foo_y_hh"},
+        ),
+        (
+            PolicyEnvironment(
+                raw_objects_tree={"foo": policy_function(leaf_name="foo")(lambda x: x)},
+                aggregation_specs_tree={},
+            ),
+            ["m", "y"],
+            ["hh"],
+            {"foo", "foo_hh"},
+        ),
+        (
+            PolicyEnvironment(
+                raw_objects_tree={},
+                aggregation_specs_tree={
+                    "foo_hh": AggregateByGroupSpec(
+                        source="foo",
+                        aggr=AggregationType.SUM,
+                    ),
+                },
+            ),
+            ["m", "y"],
+            ["hh"],
+            {"foo", "foo_hh"},
+        ),
+    ],
+)
+def test_get_top_level_namespace(
+    environment, supported_time_conversions, supported_groupings, expected
+):
+    result = _get_top_level_namespace(
+        environment=environment,
+        supported_time_conversions=supported_time_conversions,
+        supported_groupings=supported_groupings,
+    )
+    assert result == expected

tests/ttsim/test_shared.py

@@ -5,6 +5,7 @@
 from ttsim.shared import (
     create_tree_from_path_and_value,
     get_name_of_group_by_id,
+    get_re_pattern_for_time_units_and_groupings,
     insert_path_and_value,
     merge_trees,
     partition_tree_by_reference_tree,
@@ -341,3 +342,39 @@ def test_potential_target_names_from_base_name(
         )
         == expected
     )
+
+
+@pytest.mark.parametrize(
+    (
+        "func_name",
+        "supported_time_units",
+        "supported_groupings",
+        "expected_base_name",
+        "expected_time_unit",
+        "expected_aggregation",
+    ),
+    [
+        ("foo", ["m", "y"], ["hh"], "foo", None, None),
+        ("foo_m_hh", ["m", "y"], ["hh"], "foo", "m", "hh"),
+        ("foo_y_hh", ["m", "y"], ["hh"], "foo", "y", "hh"),
+        ("foo_m", ["m", "y"], ["hh"], "foo", "m", None),
+        ("foo_y", ["m", "y"], ["hh"], "foo", "y", None),
+        ("foo_hh", ["m", "y"], ["hh"], "foo", None, "hh"),
+    ],
+)
+def test_get_re_pattern_for_time_units_and_groupings(
+    func_name,
+    supported_time_units,
+    supported_groupings,
+    expected_base_name,
+    expected_time_unit,
+    expected_aggregation,
+):
+    result = get_re_pattern_for_time_units_and_groupings(
+        supported_time_units=supported_time_units,
+        supported_groupings=supported_groupings,
+    )
+    match = result.fullmatch(func_name)
+    assert match.group("base_name") == expected_base_name
+    assert match.group("time_unit") == expected_time_unit
+    assert match.group("aggregation") == expected_aggregation