refact(jetsam): simpler API no contextlib...

but cannot etablish call correctnes, relying on recetn TCs.

Author: ankostis

graphkit/base.py

@@ -1,6 +1,5 @@
 # Copyright 2016, Yahoo Inc.
 # Licensed under the terms of the Apache License, Version 2.0. See the LICENSE file associated with the project for terms.
-import contextlib
 import logging
 from collections import namedtuple
 
@@ -15,18 +14,22 @@
 log = logging.getLogger(__name__)
 
 
-@contextlib.contextmanager
-## def jetsam(locs, *salvage_vars, annotation="graphkit_jetsam", **salvage_mappings):  # bad PY2 syntax
-def jetsam(locs, *salvage_vars, **salvage_mappings):
+## def jetsam(ex, locs, *salvage_vars, annotation="graphkit_jetsam", **salvage_mappings):  # bad PY2 syntax
+def jetsam(ex, locs, *salvage_vars, **salvage_mappings):
     """
-    Debug-aid to annotate exceptions with salvaged values from wrapped functions.
+    Annotate exception with salvaged values from locals().
 
+    :param ex:
+        the exception to annotate
     :param locs:
         ``locals()`` from the context-manager's block containing vars
         to be salvaged in case of exception
 
         ATTENTION: wrapped function must finally call ``locals()``, because
         *locals* dictionary only reflects local-var changes after call.
+    :param str annotation:
+        (a kwarg not seen in the signature due to PY2 compatibility)
+        the name of the attribute to attach on the exception
     :param salvage_vars:
         local variable names to save as is in the salvaged annotations dictionary.
     :param salvage_mappings:
@@ -50,21 +53,19 @@ def jetsam(locs, *salvage_vars, **salvage_mappings):
     in case of errors::
 
 
-        with jetsam(locals(), "a", b="salvaged_b", c_var="c"):
-            try:
-                a = 1
-                b = 2
-                raise Exception()
-            finally:
-                locals()  # to update locals-dict handed to jetsam().
+        try:
+            a = 1
+            b = 2
+            raise Exception()
+        exception Exception as ex:
+            jetsam(ex, locals(), "a", b="salvaged_b", c_var="c")
 
     And then from a REPL::
 
         import sys
         sys.last_value.graphkit_jetsam
         {'a': 1, 'salvaged_b': 2, "c_var": None}
 
-
     ** Reason:**
 
     Graphs may become arbitrary deep.  Debugging such graphs is notoriously hard.
@@ -80,6 +81,7 @@ def jetsam(locs, *salvage_vars, **salvage_mappings):
     #
     annotation = salvage_mappings.pop("annotation", "graphkit_jetsam")
 
+    assert isinstance(ex, Exception), ("Bad `ex`, not an exception dict:", ex)
     assert isinstance(locs, dict), ("Bad `locs`, not a dict:", locs)
     assert all(isinstance(i, str) for i in salvage_vars), (
         "Bad `salvage_vars`!",
@@ -97,30 +99,26 @@ def jetsam(locs, *salvage_vars, **salvage_mappings):
             salvage_mappings[var] = var
 
     try:
-        yield jetsam
-    except Exception as ex_to_annotate:
-        try:
-            annotations = getattr(ex_to_annotate, annotation, None)
-            if not isinstance(annotations, dict):
-                annotations = {}
-                setattr(ex_to_annotate, annotation, annotations)
-
-            ## Salvage those asked
-            for dst_key, src in salvage_mappings.items():
-                try:
-                    salvaged_value = src(locs) if callable(src) else locs.get(src)
-                    annotations.setdefault(dst_key, salvaged_value)
-                except Exception as ex:
-                    log.warning(
-                        "Supressed error while salvaging jetsam item (%r, %r): %r"
-                        % (dst_key, src, ex)
-                    )
-        except Exception as ex:
-            log.warning(
-                "Supressed error while annotating exception: %r", ex, exc_info=1
-            )
+        annotations = getattr(ex, annotation, None)
+        if not isinstance(annotations, dict):
+            annotations = {}
+            setattr(ex, annotation, annotations)
 
-        raise  # re-raise without ex-arg, not to insert my frame
+        ## Salvage those asked
+        for dst_key, src in salvage_mappings.items():
+            try:
+                salvaged_value = src(locs) if callable(src) else locs.get(src)
+                annotations.setdefault(dst_key, salvaged_value)
+            except Exception as ex:
+                log.warning(
+                    "Supressed error while salvaging jetsam item (%r, %r): %r"
+                    % (dst_key, src, ex)
+                )
+    except Exception as ex2:
+        log.warning("Supressed error while annotating exception: %r", ex2, exc_info=1)
+        raise ex2
+
+    raise  # noqa #re-raise without ex-arg, not to insert my frame
 
 
 class Data(object):
@@ -214,23 +212,22 @@ def compute(self, inputs):
         raise NotImplementedError("Define callable of %r!" % self)
 
     def _compute(self, named_inputs, outputs=None):
-        with jetsam(
-            locals(), "outputs", "provides", "args", "results", operation="self"
-        ):
-            try:
-                provides = self.provides
-                args = [named_inputs[d] for d in self.needs]
-                results = self.compute(args)
+        try:
+            provides = self.provides
+            args = [named_inputs[d] for d in self.needs]
+            results = self.compute(args)
 
-                results = zip(provides, results)
+            results = zip(provides, results)
 
-                if outputs:
-                    outs = set(outputs)
-                    results = filter(lambda x: x[0] in outs, results)
+            if outputs:
+                outs = set(outputs)
+                results = filter(lambda x: x[0] in outs, results)
 
-                return dict(results)
-            finally:
-                locals()  # to update locals-dict handed to jetsam()
+            return dict(results)
+        except Exception as ex:
+            jetsam(
+                ex, locals(), "outputs", "provides", "args", "results", operation="self"
+            )
 
     def _after_init(self):
         """

graphkit/functional.py

@@ -14,52 +14,55 @@ def __init__(self, **kwargs):
         Operation.__init__(self, **kwargs)
 
     def _compute(self, named_inputs, outputs=None):
-        with jetsam(
-            locals(),
-            "outputs",
-            "provides",
-            "results",
-            operation="self",
-            args=lambda locs: {"args": locs.get("args"), "kwargs": locs.get("kwargs")},
-        ):
-            try:
-                args = [
-                    named_inputs[n]
-                    for n in self.needs
-                    if not isinstance(n, optional) and not isinstance(n, sideffect)
-                ]
-
-                # Find any optional inputs in named_inputs.  Get only the ones that
-                # are present there, no extra `None`s.
-                optionals = {
-                    n: named_inputs[n]
-                    for n in self.needs
-                    if isinstance(n, optional) and n in named_inputs
-                }
-
-                # Combine params and optionals into one big glob of keyword arguments.
-                kwargs = {k: v for d in (self.params, optionals) for k, v in d.items()}
-
-                # Don't expect sideffect outputs.
-                provides = [n for n in self.provides if not isinstance(n, sideffect)]
-
-                results = self.fn(*args, **kwargs)
-
-                if not provides:
-                    # All outputs were sideffects.
-                    return {}
-
-                if len(provides) == 1:
-                    results = [results]
-
-                results = zip(provides, results)
-                if outputs:
-                    outputs = set(n for n in outputs if not isinstance(n, sideffect))
-                    results = filter(lambda x: x[0] in outputs, results)
-
-                return dict(results)
-            finally:
-                locals()  # to update locals-dict handed to jetsam()
+        try:
+            args = [
+                named_inputs[n]
+                for n in self.needs
+                if not isinstance(n, optional) and not isinstance(n, sideffect)
+            ]
+
+            # Find any optional inputs in named_inputs.  Get only the ones that
+            # are present there, no extra `None`s.
+            optionals = {
+                n: named_inputs[n]
+                for n in self.needs
+                if isinstance(n, optional) and n in named_inputs
+            }
+
+            # Combine params and optionals into one big glob of keyword arguments.
+            kwargs = {k: v for d in (self.params, optionals) for k, v in d.items()}
+
+            # Don't expect sideffect outputs.
+            provides = [n for n in self.provides if not isinstance(n, sideffect)]
+
+            results = self.fn(*args, **kwargs)
+
+            if not provides:
+                # All outputs were sideffects.
+                return {}
+
+            if len(provides) == 1:
+                results = [results]
+
+            results = zip(provides, results)
+            if outputs:
+                outputs = set(n for n in outputs if not isinstance(n, sideffect))
+                results = filter(lambda x: x[0] in outputs, results)
+
+            return dict(results)
+        except Exception as ex:
+            jetsam(
+                ex,
+                locals(),
+                "outputs",
+                "provides",
+                "results",
+                operation="self",
+                args=lambda locs: {
+                    "args": locs.get("args"),
+                    "kwargs": locs.get("kwargs"),
+                },
+            )
 
     def __call__(self, *args, **kwargs):
         return self.fn(*args, **kwargs)

graphkit/network.py

@@ -259,11 +259,10 @@ def _pin_data_in_solution(self, value_name, solution, inputs, overwrites):
     def _call_operation(self, op, solution):
         # Although `plan` have added to jetsam in `compute()``,
         # add it again, in case compile()/execute is called separately.
-        with jetsam(locals(), plan="self"):
-            try:
-                return op._compute(solution)
-            finally:
-                locals()  # to update locals-dict handed to jetsam()
+        try:
+            return op._compute(solution)
+        except Exception as ex:
+            jetsam(ex, locals(), plan="self")
 
     def _execute_thread_pool_barrier_method(
         self, inputs, solution, overwrites, thread_pool_size=10
@@ -717,30 +716,29 @@ def compute(self, named_inputs, outputs, method=None, overwrites_collector=None)
 
         :returns: a dictionary of output data objects, keyed by name.
         """
-        with jetsam(locals(), "plan", "solution", "outputs", network="self"):
-            try:
-                if isinstance(outputs, str):
-                    outputs = [outputs]
-                elif not isinstance(outputs, (list, tuple)) and outputs is not None:
-                    raise ValueError(
-                        "The outputs argument must be a list, was: %s", outputs
-                    )
-
-                # Build the execution plan.
-                self.last_plan = plan = self.compile(named_inputs.keys(), outputs)
-
-                # start with fresh data solution.
-                solution = dict(named_inputs)
-
-                plan.execute(solution, overwrites_collector, method)
-
-                if outputs:
-                    # Filter outputs to just return what's requested.
-                    # Otherwise, return the whole solution as output,
-                    # including input and intermediate data nodes.
-                    # Still needed with eviction to clean isolated given inputs.
-                    solution = dict(i for i in solution.items() if i[0] in outputs)
-
-                return solution
-            finally:
-                locals()  # to update locals-dict handed to jetsam()
+        try:
+            if isinstance(outputs, str):
+                outputs = [outputs]
+            elif not isinstance(outputs, (list, tuple)) and outputs is not None:
+                raise ValueError(
+                    "The outputs argument must be a list, was: %s", outputs
+                )
+
+            # Build the execution plan.
+            self.last_plan = plan = self.compile(named_inputs.keys(), outputs)
+
+            # start with fresh data solution.
+            solution = dict(named_inputs)
+
+            plan.execute(solution, overwrites_collector, method)
+
+            if outputs:
+                # Filter outputs to just return what's requested.
+                # Otherwise, return the whole solution as output,
+                # including input and intermediate data nodes.
+                # Still needed with eviction to clean isolated given inputs.
+                solution = dict(i for i in solution.items() if i[0] in outputs)
+
+            return solution
+        except Exception as ex:
+            jetsam(ex, locals(), "plan", "solution", "outputs", network="self")