(hamilton) branch main updated: Allow inject to compose with parameterize (#1650)

[sk...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

skrawcz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hamilton.git


The following commit(s) were added to refs/heads/main by this push:
     new f15267da Allow inject to compose with parameterize (#1650)
f15267da is described below

commit f15267da5ee56bbe11d704da2bb2c4b8e8615d7c
Author: Probablism <lu...@gmail.com>
AuthorDate: Sat Sep 26 16:13:47 2026 -0400

    Allow inject to compose with parameterize (#1650)
    
    Fixes #160.
    
    This updates expander resolution so `@inject` can run before regular expanders such as `@parameterize`. That lets an injected dependency be available when the parameterized function variants are created.
    
    ## Changes
    
    - Add a pre-expander pass for expanders that need to run before regular node expansion.
    - Mark `@inject` as a pre-expander while keeping the existing single regular expander restriction.
    - Validate that injected target parameters do not overlap with parameters consumed by regular parameterization.
    - Add regression coverage for both decorator orders and the overlapping-parameter failure case.
    
    ## How I tested this
    
    - `tests/function_modifiers/test_expanders.py` passed locally.
    - The main package test suite passed locally: 1446 passed, 5 skipped, 2 xpassed.
    - Static lint/format checks passed locally.
    - Fork GitHub Actions Unit Tests passed on `fix-inject-parameterize` across Python 3.10 through 3.14.
    
    ## Notes
    
    - No user-facing documentation changes were needed; this is a decorator composition bug fix.
    
    ## Checklist
    
    - [x] PR has an informative and human-readable title (this will be pulled into the release notes)
    - [x] Changes are limited to a single goal (no scope creep)
    - [x] Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
    - [x] Any _change_ in functionality is tested
    - [x] New functions are documented (with a description, list of inputs, and expected output)
    - [x] Placeholder code is flagged / future TODOs are captured in comments
    - [x] Project documentation has been updated if adding/changing functionality.
---
 hamilton/function_modifiers/base.py        | 52 +++++++++++++++++++++++++++---
 hamilton/function_modifiers/expanders.py   |  4 +++
 tests/function_modifiers/test_expanders.py | 44 +++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 4 deletions(-)

diff --git a/hamilton/function_modifiers/base.py b/hamilton/function_modifiers/base.py
index 10aa9b39..b3f9dcdb 100644
--- a/hamilton/function_modifiers/base.py
+++ b/hamilton/function_modifiers/base.py
@@ -337,6 +337,10 @@ class NodeExpander(SubDAGModifier):
 
     EXPAND_NODES = "expand_nodes"
 
+    @classmethod
+    def runs_before_regular_expanders(cls) -> bool:
+        return False
+
     def transform_dag(
         self, nodes: Collection[node.Node], config: dict[str, Any], fn: Callable
     ) -> Collection[node.Node]:
@@ -794,6 +798,35 @@ def _resolve_nodes_error(fn: Callable) -> str:
     return f"Exception occurred while compiling function: {fn.__name__} to nodes"
 
 
+def _parameterized_targets(expander: NodeExpander) -> set[str]:
+    """Return original function parameters replaced by parameterization-style expanders."""
+    parameterization = getattr(expander, "parameterization", {})
+    targets = set()
+    for mapping in parameterization.values():
+        if isinstance(mapping, tuple):
+            mapping = mapping[0]
+        if isinstance(mapping, dict):
+            targets.update(mapping)
+    return targets
+
+
+def _validate_expander_composition(
+    pre_expanders: list[NodeExpander], regular_expanders: list[NodeExpander]
+) -> None:
+    injected_targets = set().union(
+        *(_parameterized_targets(expander) for expander in pre_expanders)
+    )
+    regular_targets = set().union(
+        *(_parameterized_targets(expander) for expander in regular_expanders)
+    )
+    overlap = injected_targets & regular_targets
+    if overlap:
+        raise InvalidDecoratorException(
+            "Cannot combine @inject and @parameterize replacements for the same parameter(s): "
+            f"{', '.join(sorted(overlap))}"
+        )
+
+
 def resolve_nodes(fn: Callable, config: dict[str, Any]) -> Collection[node.Node]:
     """Gets a list of nodes from a function. This is meant to be an abstraction between the node
     and the function that it implements. This will end up coordinating with the decorators we build
@@ -810,8 +843,8 @@ def resolve_nodes(fn: Callable, config: dict[str, Any]) -> Collection[node.Node]
     -- this is determined in the node creator class. Apply that to get
     the initial node.
 
-    3. If there is a list of node expanders, apply them. Otherwise apply the default
-    node expander This must be a list of length one. This gives out a list of nodes.
+    3. If there is a list of node expanders, apply pre-expanders before at most one regular
+    expander. This gives out a list of nodes.
 
     4. If there is a node transformer, apply that. Note that the node transformer
     gets applied individually to just the sink nodes in the subdag. It subclasses
@@ -837,8 +870,19 @@ def resolve_nodes(fn: Callable, config: dict[str, Any]) -> Collection[node.Node]
         for node_injector in node_injectors:
             nodes = node_injector.transform_dag(nodes, filter_config(config, node_injector), fn)
         node_expanders = function_decorators[NodeExpander.get_lifecycle_name()]
-        if len(node_expanders) > 0:
-            (node_expander,) = node_expanders
+        pre_expanders = [
+            expander for expander in node_expanders if expander.runs_before_regular_expanders()
+        ]
+        regular_expanders = [
+            expander for expander in node_expanders if not expander.runs_before_regular_expanders()
+        ]
+        if len(regular_expanders) > 1:
+            raise InvalidDecoratorException(
+                "Cannot combine multiple regular node expanders on one function: "
+                f"{', '.join(expander.name for expander in regular_expanders)}"
+            )
+        _validate_expander_composition(pre_expanders, regular_expanders)
+        for node_expander in pre_expanders + regular_expanders:
             nodes = node_expander.transform_dag(nodes, filter_config(config, node_expander), fn)
         node_transformers = function_decorators[NodeTransformer.get_lifecycle_name()]
         for dag_modifier in node_transformers:
diff --git a/hamilton/function_modifiers/expanders.py b/hamilton/function_modifiers/expanders.py
index a4141d12..c25266b9 100644
--- a/hamilton/function_modifiers/expanders.py
+++ b/hamilton/function_modifiers/expanders.py
@@ -1231,3 +1231,7 @@ class inject(parameterize):
             This is the same as the input mapping in `@parameterize`.
         """
         super(inject, self).__init__(**{parameterize.PLACEHOLDER_PARAM_NAME: key_mapping})
+
+    @classmethod
+    def runs_before_regular_expanders(cls) -> bool:
+        return True
diff --git a/tests/function_modifiers/test_expanders.py b/tests/function_modifiers/test_expanders.py
index d9c37850..a5b4edfe 100644
--- a/tests/function_modifiers/test_expanders.py
+++ b/tests/function_modifiers/test_expanders.py
@@ -1088,6 +1088,50 @@ def test_inject_multiple_things():
     assert node_(int_value_not_injected=8, three_source=3, six_source=6) == 8 * (8 + 1) // 2
 
 
+def test_inject_above_parameterize_resolves_nodes():
+    @function_modifiers.inject(params=source("my_func__params"))
+    @function_modifiers.parameterize(
+        my_func_a={"date_range": source("my_func_a_date_range")},
+        my_func_b={"date_range": source("my_func_b_date_range")},
+    )
+    def my_func(date_range: tuple[int, int], params: int) -> int:
+        return params + date_range[0] + date_range[1]
+
+    nodes = {node_.name: node_ for node_ in base.resolve_nodes(my_func, {})}
+
+    assert set(nodes) == {"my_func_a", "my_func_b"}
+    assert nodes["my_func_a"].callable(my_func_a_date_range=(2, 10), my_func__params=1) == 13
+    assert nodes["my_func_b"].callable(my_func_b_date_range=(3, 10), my_func__params=1) == 14
+
+
+def test_parameterize_above_inject_resolves_nodes():
+    @function_modifiers.parameterize(
+        my_func_a={"date_range": source("my_func_a_date_range")},
+        my_func_b={"date_range": source("my_func_b_date_range")},
+    )
+    @function_modifiers.inject(params=source("my_func__params"))
+    def my_func(date_range: tuple[int, int], params: int) -> int:
+        return params + date_range[0] + date_range[1]
+
+    nodes = {node_.name: node_ for node_ in base.resolve_nodes(my_func, {})}
+
+    assert set(nodes) == {"my_func_a", "my_func_b"}
+    assert nodes["my_func_a"].callable(my_func_a_date_range=(2, 10), my_func__params=1) == 13
+    assert nodes["my_func_b"].callable(my_func_b_date_range=(3, 10), my_func__params=1) == 14
+
+
+def test_inject_parameterize_overlap_fails():
+    @function_modifiers.inject(date_range=source("injected_date_range"))
+    @function_modifiers.parameterize(
+        my_func_a={"date_range": source("my_func_a_date_range")},
+    )
+    def my_func(date_range: tuple[int, int], params: int) -> int:
+        return params + date_range[0] + date_range[1]
+
+    with pytest.raises(base.InvalidDecoratorException, match="date_range"):
+        base.resolve_nodes(my_func, {})
+
+
 @pytest.mark.parametrize(
     ("annotated_type", "cls", "expected"),
     [