Fix bug in LabelledPipeline.fix

Author: Matthieu Ancellin

labelled_functions/pipeline.py

@@ -2,7 +2,7 @@
 # coding: utf-8
 
 from typing import Set
-from collections import namedtuple
+from collections import namedtuple, defaultdict
 from toolz.itertoolz import groupby
 from toolz.dicttoolz import merge, keyfilter
 
@@ -36,7 +36,8 @@ def __init__(self,
         if any(f.output_names is Unknown for f in self.funcs):
             raise AttributeError("Cannot build a pipeline with a function whose outputs are unknown.")
 
-        pipe_inputs, sub_default_values, pipe_outputs, *_ = self._graph()
+        self._graph_data = self._graph()
+        pipe_inputs, sub_default_values, pipe_outputs, *_ = self._graph_data
 
         if default_values is None:
             sub_default_values = sub_default_values
@@ -93,15 +94,22 @@ def __ror__(self, other):
         else:
             return NotImplemented
 
+    def _which_input_is_used_by_this_function(self):
+        _, _, _, _, _, edges = self._graph_data
+        inputs_used_by = defaultdict(set)
+        for e in edges:
+            if e.start is None:
+                inputs_used_by[e.end].add(e.label)
+        return inputs_used_by
+
     def fix(self, **names_to_fix):
+        inputs_used_by = self._which_input_is_used_by_this_function()
+
         fixed_funcs = []
         for f in self.funcs:
-            fixable_names = {n: v for n, v in names_to_fix.items() if n in f.input_names}
+            fixable_names = {k: v for k, v in names_to_fix.items() if k in inputs_used_by[f.name]}
             if len(fixable_names) > 0:
                 fixed_funcs.append(f.fix(**fixable_names))
-
-                # Remove inputs that have been fixed
-                names_to_fix = {n: v for n, v in names_to_fix.items() if n not in fixable_names}
             else:
                 fixed_funcs.append(f)
 

test/test_pipeline.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from labelled_functions.labels import LabelledFunction
+from labelled_functions.labels import LabelledFunction, label
 from labelled_functions.pipeline import LabelledPipeline, pipeline, compose
 from labelled_functions.special_functions import let, relabel, show
 from labelled_functions.decorators import keeping_inputs
@@ -66,9 +66,6 @@ def g(b, c=3):
     assert pipe(a=1) == {'output': 1}
     assert pipe(a=1, c=0) == {'output': -2}
 
-    pipe = pipeline([f, g]).fix(a=1)
-    assert pipe() == {'output': 1}
-
     ###
     pipe = pipeline([cube], default_values={'x': 1.0})
     assert pipe()['volume'] == 1.0
@@ -79,6 +76,19 @@ def g(b, c=3):
     assert pipe(x=10.0)['volume'] == 1000.0
 
 
+def test_fix():
+    a, b = 1, 2
+    f = label(lambda x, y: x + y, name="foo", output_names=['z'])
+    g = label(lambda x, z: x * z, name="bar", output_names=['u'])
+    h = label(lambda u: u**2, name="baz", output_names=['w'])
+
+    pipe = pipeline([f, g, h]).fix(y=b)
+    assert pipe(x=a) == {'w': (a*(a+b))**2}
+
+    pipe = pipeline([f, g, h]).fix(x=a)
+    assert pipe(y=b) == {'w': (a*(a+b))**2}
+
+
 def test_let():
     l = let(x=1.0)
     assert l.input_names == []