fix: DAG rendering from code package (Netflix#290)

* inherit DAGNode from metaflow client instead of reimplementing it.

* inherit most of FlowGraph functionality from metaflow client as well, overwriting only when necessary.

* draft unit test for legacy dag parsing

* first passing unit test for custom flowgraph

* add more unit tests to cover custom flowgraph behaviour

* revert client FlowGraph inheritance

* codestyles

* fix compatibility with imported DAGNode

Co-authored-by: Brendan Gibson <brendan@outerbounds.co>

Author: saikonen

services/ui_backend_service/data/cache/custom_flowgraph.py

@@ -1,110 +1,9 @@
 import ast
+from metaflow.graph import deindent_docstring, DAGNode
 
-
-class DAGNode(object):
-    def __init__(self, func_ast, decos, doc):
-        self.name = func_ast.name
-        self.func_lineno = func_ast.lineno
-        self.decorators = decos
-        self.doc = doc.rstrip()
-
-        # these attributes are populated by _parse
-        self.tail_next_lineno = 0
-        self.type = None
-        self.out_funcs = []
-        self.has_tail_next = False
-        self.invalid_tail_next = False
-        self.num_args = 0
-        self.condition = None
-        self.foreach_param = None
-        self._parse(func_ast)
-
-        # these attributes are populated by _traverse_graph
-        self.in_funcs = set()
-        self.split_parents = []
-        self.matching_join = None
-
-        # these attributes are populated by _postprocess
-        self.is_inside_foreach = False
-
-    def _expr_str(self, expr):
-        return '%s.%s' % (expr.value.id, expr.attr)
-
-    def _parse(self, func_ast):
-
-        self.num_args = len(func_ast.args.args)
-        tail = func_ast.body[-1]
-
-        # end doesn't need a transition
-        if self.name == 'end':
-            # TYPE: end
-            self.type = 'end'
-
-        # ensure that the tail an expression
-        if not isinstance(tail, ast.Expr):
-            return
-
-        # determine the type of self.next transition
-        try:
-            if not self._expr_str(tail.value.func) == 'self.next':
-                return
-
-            self.has_tail_next = True
-            self.invalid_tail_next = True
-            self.tail_next_lineno = tail.lineno
-            self.out_funcs = [e.attr for e in tail.value.args]
-            keywords = dict((k.arg, k.value.s) for k in tail.value.keywords)
-
-            if len(keywords) == 1:
-                if 'foreach' in keywords:
-                    # TYPE: foreach
-                    self.type = 'foreach'
-                    if len(self.out_funcs) == 1:
-                        self.foreach_param = keywords['foreach']
-                        self.invalid_tail_next = False
-                elif 'condition' in keywords:
-                    # TYPE: split-or
-                    self.type = 'split-or'
-                    if len(self.out_funcs) == 2:
-                        self.condition = keywords['condition']
-                        self.invalid_tail_next = False
-            elif len(keywords) == 0:
-                if len(self.out_funcs) > 1:
-                    # TYPE: split-and
-                    self.type = 'split-and'
-                    self.invalid_tail_next = False
-                elif len(self.out_funcs) == 1:
-                    # TYPE: linear
-                    if self.num_args > 1:
-                        self.type = 'join'
-                    else:
-                        self.type = 'linear'
-                    self.invalid_tail_next = False
-
-        except AttributeError:
-            return
-
-    def __str__(self):
-        return """
-            *[{0.name} {0.type} (line {0.func_lineno})]*
-            in_funcs={in_funcs}
-            split_parents={parents}
-            matching_join={matching_join}
-            is_inside_foreach={is_inside_foreach}
-            decorators={decos}
-            num_args={0.num_args}
-            has_tail_next={0.has_tail_next} (line {0.tail_next_lineno})
-            invalid_tail_next={0.invalid_tail_next}
-            condition={0.condition}
-            foreach_param={0.foreach_param}
-            -> {out}"""\
-            .format(self,
-                    matching_join=self.matching_join and '[%s]' % self.matching_join,
-                    is_inside_foreach=self.is_inside_foreach,
-                    in_funcs=', '.join('[%s]' % x for x in self.in_funcs),
-                    parents=', '.join('[%s]' % x for x in self.split_parents),
-                    decos=' | '.join(map(str, self.decorators)),
-                    out=', '.join('[%s]' % x for x in self.out_funcs))
+# NOTE: This is a custom implementation of the FlowGraph class from the Metaflow client
+# which can parse a graph out of a flow_name and a source code string, instead of relying on
+# importing the source code as a module.
 
 
 class StepVisitor(ast.NodeVisitor):
@@ -141,7 +40,7 @@ def _flow(n):
         [root] = list(filter(_flow, ast.parse(source).body))
         self.name = root.name
         doc = ast.get_docstring(root)
-        self.doc = doc if doc else ''
+        self.doc = deindent_docstring(doc) if doc else ''
         nodes = {}
         StepVisitor(nodes).visit(root)
         return nodes
@@ -151,20 +50,18 @@ def _postprocess(self):
         # has is_inside_foreach=True *unless* all of those foreaches
         # are joined by the node
         for node in self.nodes.values():
-            foreaches = [p for p in node.split_parents
-                         if self.nodes[p].type == 'foreach']
-            if [f for f in foreaches
-                    if self.nodes[f].matching_join != node.name]:
+            foreaches = [
+                p for p in node.split_parents if self.nodes[p].type == "foreach"
+            ]
+            if [f for f in foreaches if self.nodes[f].matching_join != node.name]:
                 node.is_inside_foreach = True
 
     def _traverse_graph(self):
-
         def traverse(node, seen, split_parents):
-
-            if node.type in ('split-or', 'split-and', 'foreach'):
+            if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
                 split_parents = split_parents + [node.name]
-            elif node.type == 'join':
+            elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
                     self[split_parents[-1]].matching_join = node.name
@@ -182,8 +79,8 @@ def traverse(node, seen, split_parents):
                         child.in_funcs.add(node.name)
                         traverse(child, seen + [n], split_parents)
 
-        if 'start' in self:
-            traverse(self['start'], [], [])
+        if "start" in self:
+            traverse(self["start"], [], [])
 
         # fix the order of in_funcs
         for node in self.nodes.values():
@@ -199,8 +96,9 @@ def __iter__(self):
         return iter(self.nodes.values())
 
     def __str__(self):
-        return '\n'.join(str(n) for _, n in sorted((n.func_lineno, n)
-                                                   for n in self.nodes.values()))
+        return "\n".join(
+            str(n) for _, n in sorted((n.func_lineno, n) for n in self.nodes.values())
+        )
 
     def output_steps(self):