AST-only version of graph.FlowGraph (#249)

* allow graph.FlowGraph to be used just with the source code

this avoids importing the flow module and execution of any user-defined code

* parse @Deco() style decorators correctly in AST-only FlowGraph

Co-authored-by: Ville Tuulos <[email protected]>

Author: tuulos

metaflow/graph.py

@@ -2,49 +2,12 @@
 import ast
 import re
 
-def deindent_docstring(doc):
-    if doc:
-        # Find the indent to remove from the doctring. We consider the following possibilities:
-        # Option 1:
-        #  """This is the first line
-        #    This is the second line
-        #  """
-        # Option 2:
-        #  """
-        # This is the first line
-        # This is the second line
-        # """
-        # Option 3:
-        #  """
-        #     This is the first line
-        #     This is the second line
-        #  """
-        #
-        # In all cases, we can find the indent to remove by doing the following:
-        #  - Check the first non-empty line, if it has an indent, use that as the base indent
-        #  - If it does not have an indent and there is a second line, check the indent of the
-        #    second line and use that
-        saw_first_line = False
-        matched_indent = None
-        for line in doc.splitlines():
-            if line:
-                matched_indent = re.match('[\t ]+', line)
-                if matched_indent is not None or saw_first_line:
-                    break
-                saw_first_line = True
-        if matched_indent:
-            return re.sub(r'\n' + matched_indent.group(), '\n', doc).strip()
-        else:
-            return doc
-    else:
-        return ''
-
 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 = deindent_docstring(doc)
+        self.doc = doc.rstrip()
 
         # these attributes are populated by _parse
         self.tail_next_lineno = 0
@@ -148,32 +111,46 @@ def __str__(self):
 
 class StepVisitor(ast.NodeVisitor):
 
-    def __init__(self, nodes, flow):
+    def __init__(self, nodes):
         self.nodes = nodes
-        self.flow = flow
         super(StepVisitor, self).__init__()
 
     def visit_FunctionDef(self, node):
-        func = getattr(self.flow, node.name)
-        if hasattr(func, 'is_step'):
-            self.nodes[node.name] = DAGNode(node, func.decorators, func.__doc__)
+        decos = [d.func.id if isinstance(d, ast.Call) else d.id
+                 for d in node.decorator_list]
+        if 'step' in decos:
+            doc = ast.get_docstring(node)
+            self.nodes[node.name] = DAGNode(node, decos, doc if doc else '')
 
 class FlowGraph(object):
 
-    def __init__(self, flow):
-        self.name = flow.__name__
-        self.nodes = self._create_nodes(flow)
-        self.doc = deindent_docstring(flow.__doc__)
+    def __init__(self, flow=None, source=None, name=None):
+        if flow:
+            module = __import__(flow.__module__)
+            source = inspect.getsource(module)
+            self.name = flow.__name__
+        else:
+            self.name = name
+
+        self.nodes = self._create_nodes(source)
         self._traverse_graph()
         self._postprocess()
 
-    def _create_nodes(self, flow):
-        module = __import__(flow.__module__)
-        tree = ast.parse(inspect.getsource(module)).body
-        root = [n for n in tree\
-                if isinstance(n, ast.ClassDef) and n.name == self.name][0]
+    def _create_nodes(self, source):
+        def _flow(n):
+            if isinstance(n, ast.ClassDef):
+                bases = [b.id for b in n.bases]
+                if 'FlowSpec' in bases:
+                    return self.name is None or n.name == self.name
+
+        # NOTE: this will fail if a file has multiple FlowSpec classes
+        # and no name is specified
+        [root] = list(filter(_flow, ast.parse(source).body))
+        self.name = root.name
+        doc = ast.get_docstring(root)
+        self.doc = doc if doc else ''
         nodes = {}
-        StepVisitor(nodes, flow).visit(root)
+        StepVisitor(nodes).visit(root)
         return nodes
 
     def _postprocess(self):