Make top-level namespace include potential derived functions

src/ttsim/compute_taxes_and_transfers.py

@@ -31,8 +31,10 @@
     format_list_linewise,
     get_name_of_group_by_id,
     get_names_of_arguments_without_defaults,
+    get_re_pattern_for_time_units_and_groupings,
     merge_trees,
     partition_by_reference_dict,
+    potential_target_names_from_base_name,
 )
 from ttsim.typing import (
     check_series_has_expected_type,
@@ -88,9 +90,7 @@ def compute_taxes_and_transfers(
     _fail_if_environment_not_valid(environment)
 
     # Transform functions tree to qualified names dict with qualified arguments
-    top_level_namespace = set(environment.raw_objects_tree.keys()) | set(
-        environment.aggregation_specs_tree.keys()
-    )
+    top_level_namespace = _get_top_level_namespace(environment)
     functions = dt.functions_without_tree_logic(
         functions=environment.functions_tree, top_level_namespace=top_level_namespace
     )
@@ -166,6 +166,51 @@ def compute_taxes_and_transfers(
     return result_tree
 
 
+def _get_top_level_namespace(
+    environment: PolicyEnvironment,
+    supported_time_conversions: list[str],
+    supported_groupings: list[str],
+) -> set[str]:
+    """Get the top level namespace.
+
+    Parameters
+    ----------
+    environment:
+        The policy environment.
+
+    Returns
+    -------
+    top_level_namespace:
+        The top level namespace.
+    """
+    names_from_environment = set(environment.raw_objects_tree.keys()) | set(
+        environment.aggregation_specs_tree.keys()
+    )
+    potential_function_names = set()
+    re_pattern = get_re_pattern_for_time_units_and_groupings(
+        supported_groupings=supported_groupings,
+        supported_time_units=supported_time_conversions,
+    )
+
+    for element in names_from_environment:
+        if match := re_pattern.fullmatch(element):
+            function_base_name = match.group("base_name")
+            create_conversions_for_time_units = bool(match.group("time_unit"))
+        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
+
+
 def _convert_data_to_correct_types(
     data: QualNameDataDict, functions_overridden: QualNameTTSIMObjectDict
 ) -> QualNameDataDict:

src/ttsim/shared.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import itertools
 import re
 import textwrap
 from typing import TYPE_CHECKING, Any, TypeVar
@@ -36,6 +37,95 @@ def validate_date_range(start: datetime.date, end: datetime.date):
         raise ValueError(f"The start date {start} must be before the end date {end}.")
 
 
+def get_re_pattern_for_time_units_and_groupings(
+    supported_groupings: list[str], supported_time_units: list[str]
+) -> re.Pattern:
+    """Get a regex pattern for time units and 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
+    ----------
+    supported_groupings
+        The supported groupings.
+    supported_time_units
+        The supported time units.
+
+    Returns
+    -------
+    pattern
+        The regex pattern.
+    """
+    units = "".join(supported_time_units)
+    groupings = "|".join(supported_groupings)
+    return re.compile(
+        f"(?P<base_name>.*?)"
+        f"(?:_(?P<time_unit>[{units}]))?"
+        f"(?:_(?P<aggregation>{groupings}))?"
+        f"$"
+    )
+
+
+def potential_target_names_from_base_name(
+    base_name: str,
+    supported_time_conversions: list[str],
+    supported_groupings: list[str],
+    create_conversions_for_time_units: bool,
+) -> set[str]:
+    """Get all derived function names given a base function name.
+
+    Examples
+    --------
+    >>> potential_target_names_from_base_name(
+        base_name="income",
+        supported_time_conversions=["y", "m"],
+        supported_groupings=["hh"],
+        create_conversions_for_time_units=True,
+    )
+    {'income_m', 'income_y', 'income_hh_y', 'income_hh_m'}
+
+    >>> potential_target_names_from_base_name(
+        base_name="claims_benefits",
+        supported_time_conversions=["y", "m"],
+        supported_groupings=["hh"],
+        create_conversions_for_time_units=False,
+    )
+    {'claims_benefits_hh'}
+
+    Parameters
+    ----------
+    base_name
+        The base function name.
+    supported_time_conversions
+        The supported time conversions.
+    supported_groupings
+        The supported groupings.
+    create_conversions_for_time_units
+        Whether to create conversions for time units.
+
+    Returns
+    -------
+    The names of all potential targets based on the base name.
+    """
+    result = set()
+    if create_conversions_for_time_units:
+        for time_unit in supported_time_conversions:
+            result.add(f"{base_name}_{time_unit}")
+        for time_unit, aggregation in itertools.product(
+            supported_time_conversions, supported_groupings
+        ):
+            result.add(f"{base_name}_{time_unit}_{aggregation}")
+    else:
+        result.add(base_name)
+        for aggregation in supported_groupings:
+            result.add(f"{base_name}_{aggregation}")
+    return result
+
+
 class KeyErrorMessage(str):
     """Subclass str to allow for line breaks in KeyError messages."""
 

src/ttsim/time_conversion.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import inspect
-import re
 from typing import TYPE_CHECKING
 
 import dags.tree as dt
 from dags import rename_arguments
 
 from _gettsim.config import SUPPORTED_GROUPINGS
 from ttsim.function_types import DerivedTimeConversionFunction, PolicyFunction
+from ttsim.shared import get_re_pattern_for_time_units_and_groupings
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -447,25 +447,28 @@ def _create_time_conversion_functions(
     name: str, func: PolicyFunction | None = None
 ) -> dict[str, DerivedTimeConversionFunction]:
     result: dict[str, DerivedTimeConversionFunction] = {}
-
     all_time_units = list(_TIME_UNITS)
 
-    units = "".join(all_time_units)
-    groupings = "|".join([f"_{grouping}" for grouping in SUPPORTED_GROUPINGS])
-    function_with_time_unit = re.compile(
-        f"(?P<base_name>.*_)(?P<time_unit>[{units}])(?P<aggregation>{groupings})?"
+    time_unit_pattern = get_re_pattern_for_time_units_and_groupings(
+        supported_groupings=SUPPORTED_GROUPINGS,
+        supported_time_units=all_time_units,
     )
-    match = function_with_time_unit.fullmatch(name)
-    dependencies = set(inspect.signature(func).parameters) if func else set()
 
-    if match:
-        base_name = match.group("base_name")
-        time_unit = match.group("time_unit")
-        aggregation = match.group("aggregation") or ""
+    match = time_unit_pattern.fullmatch(name)
+    base_name = match.group("base_name")
+    time_unit = match.group("time_unit") or ""
+    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