Ignore partialled arguments of functions. (#4)

Author: janosg

src/dags/dag.py

@@ -1,3 +1,4 @@
+import functools
 import inspect
 import textwrap
 
@@ -24,7 +25,7 @@ def concatenate_functions(
     Functions that are not required to produce the targets will simply be ignored.
 
     The arguments of the combined function are all arguments of relevant functions
-    that are not themselves function names.
+    that are not themselves function names, in alphabetical order.
 
     Args:
         functions (dict or list): Dict or list of functions. If a list, the function
@@ -152,14 +153,22 @@ def _create_complete_dag(functions):
 
     """
     functions_arguments_dict = {
-        name: list(inspect.signature(function).parameters)
-        for name, function in functions.items()
+        name: _get_free_arguments(function) for name, function in functions.items()
     }
     dag = nx.DiGraph(functions_arguments_dict).reverse()
 
     return dag
 
 
+def _get_free_arguments(func):
+    arguments = list(inspect.signature(func).parameters)
+    if isinstance(func, functools.partial):
+        non_free = set(func.args) | set(func.keywords)
+        arguments = [arg for arg in arguments if arg not in non_free]
+
+    return arguments
+
+
 def _limit_dag_to_targets_and_their_ancestors(dag, targets):
     """Limit DAG to targets and their ancestors.
 
@@ -216,7 +225,7 @@ def _create_execution_info(functions, dag):
     out = {}
     for node in nx.topological_sort(dag):
         if node in functions:
-            arguments = list(inspect.signature(functions[node]).parameters)
+            arguments = _get_free_arguments(functions[node])
             info = {}
             info["func"] = functions[node]
             info["arguments"] = arguments

tests/test_dag.py

@@ -1,4 +1,5 @@
 import inspect
+from functools import partial
 
 import pytest
 from dags.dag import concatenate_functions
@@ -119,3 +120,19 @@ def test_concatenate_functions_with_aggregation_via_or():
         aggregator=lambda a, b: a or b,
     )
     assert aggregated()
+
+
+def test_partialled_argument_is_ignored():
+    def f(a, b):
+        return a + b
+
+    def g(f, c):
+        return f + c
+
+    concatenated = concatenate_functions(
+        functions={"f": partial(f, b=2), "g": g},
+        targets="g",
+    )
+
+    assert list(inspect.signature(concatenated).parameters) == ["a", "c"]
+    assert concatenated(1, 3) == 6