Allow getting full context in partial outputs

Author: csala

mlblocks/mlpipeline.py

@@ -4,6 +4,7 @@
 
 import json
 import logging
+import re
 from collections import Counter, OrderedDict, defaultdict
 from copy import deepcopy
 
@@ -198,12 +199,15 @@ def __init__(self, pipeline=None, primitives=None, init_params=None,
         if hyperparameters:
             self.set_hyperparameters(hyperparameters)
 
+        self._re_block_name = re.compile(r'(^[^#]+#\d+)(\..*)?')
+
     def _get_str_output(self, output):
         """Get the outputs that correspond to the str specification."""
         if output in self.outputs:
             return self.outputs[output]
         elif output in self.blocks:
-            return self._get_block_outputs(output)
+            return [{'name': output, 'variable': output}]
+            # return self._get_block_outputs(output)
         elif '.' in output:
             block_name, variable_name = output.rsplit('.', 1)
             block = self.blocks.get(block_name)
@@ -260,11 +264,11 @@ def get_outputs(self, outputs='default'):
 
         computed = list()
         for output in outputs:
+            if isinstance(output, int):
+                output = self._get_block_name(output)
+
             if isinstance(output, str):
                 computed.extend(self._get_str_output(output))
-            elif isinstance(output, int):
-                block_name = self._get_block_name(output)
-                computed.extend(self._get_block_outputs(block_name))
             else:
                 raise TypeError('Output Specification can only be str or int')
 
@@ -316,6 +320,18 @@ def get_output_variables(self, outputs='default'):
         outputs = self.get_outputs(outputs)
         return [output['variable'] for output in outputs]
 
+    def _extract_block_name(self, variable_name):
+        return self._re_block_name.search(variable_name).group(1)
+
+    def _prepare_outputs(self, outputs):
+        output_variables = self.get_output_variables(outputs)
+        outputs = output_variables.copy()
+        output_blocks = {
+            self._extract_block_name(variable)
+            for variable in output_variables
+        }
+        return output_variables, outputs, output_blocks
+
     @staticmethod
     def _flatten_dict(hyperparameters):
         return {
@@ -519,13 +535,11 @@ def _extract_outputs(self, block_name, outputs, block_outputs):
 
         return output_dict
 
-    def _update_outputs(self, block_name, output_variables, outputs, outputs_dict):
+    def _update_outputs(self, variable_name, output_variables, outputs, value):
         """Set the requested block outputs into the outputs list in the right place."""
-        for key, value in outputs_dict.items():
-            variable_name = '{}.{}'.format(block_name, key)
-            if variable_name in output_variables:
-                index = output_variables.index(variable_name)
-                outputs[index] = deepcopy(value)
+        if variable_name in output_variables:
+            index = output_variables.index(variable_name)
+            outputs[index] = deepcopy(value)
 
     def _fit_block(self, block, block_name, context):
         """Get the block args from the context and fit the block."""
@@ -554,7 +568,12 @@ def _produce_block(self, block, block_name, context, output_variables, outputs):
             context.update(outputs_dict)
 
             if output_variables:
-                self._update_outputs(block_name, output_variables, outputs, outputs_dict)
+                if block_name in output_variables:
+                    self._update_outputs(block_name, output_variables, outputs, context)
+                else:
+                    for key, value in outputs_dict.items():
+                        variable_name = '{}.{}'.format(block_name, key)
+                        self._update_outputs(variable_name, output_variables, outputs, value)
 
         except Exception:
             if self.verbose:
@@ -609,17 +628,12 @@ def fit(self, X=None, y=None, output_=None, start_=None, **kwargs):
         if y is not None:
             context['y'] = y
 
-        if isinstance(output_, str):
-            output_variables = self.get_output_variables(output_)
-            outputs = output_variables.copy()
-            output_blocks = {variable.rsplit('.', 1)[0] for variable in output_variables}
-        else:
+        if output_ is None:
             output_variables = None
             outputs = None
             output_blocks = set()
-
-        if isinstance(output_, int):
-            output_ = self._get_block_name(output_)
+        else:
+            output_variables, outputs, output_blocks = self._prepare_outputs(output_)
 
         if isinstance(start_, int):
             start_ = self._get_block_name(start_)
@@ -634,15 +648,12 @@ def fit(self, X=None, y=None, output_=None, start_=None, **kwargs):
 
             self._fit_block(block, block_name, context)
 
-            last_block = block_name != self._last_block_name
-            if last_block or (block_name == output_) or (block_name in output_blocks):
+            if (block_name != self._last_block_name) or (block_name in output_blocks):
                 self._produce_block(block, block_name, context, output_variables, outputs)
 
                 # We already captured the output from this block
                 if block_name in output_blocks:
                     output_blocks.remove(block_name)
-                elif block_name == output_:
-                    return context
 
             # If there was an output_ but there are no pending
             # outputs we are done.
@@ -695,9 +706,7 @@ def predict(self, X=None, output_='default', start_=None, **kwargs):
         if X is not None:
             context['X'] = X
 
-        output_variables = self.get_output_variables(output_)
-        outputs = output_variables.copy()
-        output_blocks = {variable.rsplit('.', 1)[0] for variable in output_variables}
+        output_variables, outputs, output_blocks = self._prepare_outputs(output_)
 
         if isinstance(start_, int):
             start_ = self._get_block_name(start_)

tests/features/test_partial_outputs.py

@@ -70,13 +70,14 @@ def test_fit_output(self):
         y = np.array([
             0, 0, 0, 0, 1
         ])
+        context = {'X': X, 'y': y}
 
         almost_equal(named_out, y)
         assert len(list_out) == 2
         almost_equal(list_out[0], y)
-        almost_equal(list_out[1], X)
-        almost_equal(X, int_out)
-        almost_equal(X, str_out)
+        almost_equal(list_out[1], context)
+        almost_equal(context, int_out)
+        almost_equal(context, str_out)
         almost_equal(X, str_out_variable)
         assert no_output is None