ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/graph.py

@@ -2,6 +2,11 @@ import inspect
 import ast
 import re
 
+from itertools import chain
+
+
+from .util import to_pod
+
 
 def deindent_docstring(doc):
     if doc:
@@ -42,10 +47,17 @@ def deindent_docstring(doc):
 
 
 class DAGNode(object):
-    def __init__(self, func_ast, decos, doc):
+    def __init__(
+        self, func_ast, decos, wrappers, config_decorators, doc, source_file, lineno
+    ):
         self.name = func_ast.name
-        self.func_lineno = func_ast.lineno
+        self.source_file = source_file
+        # lineno is the start line of decorators in source_file
+        # func_ast.lineno is lines from decorators start to def of function
+        self.func_lineno = lineno + func_ast.lineno - 1
         self.decorators = decos
+        self.wrappers = wrappers
+        self.config_decorators = config_decorators
         self.doc = deindent_docstring(doc)
         self.parallel_step = any(getattr(deco, "IS_PARALLEL", False) for deco in decos)
 
@@ -56,14 +68,17 @@ class DAGNode(object):
         self.has_tail_next = False
         self.invalid_tail_next = False
         self.num_args = 0
+        self.switch_cases = {}
+        self.condition = None
         self.foreach_param = None
         self.num_parallel = 0
         self.parallel_foreach = False
-        self._parse(func_ast)
+        self._parse(func_ast, lineno)
 
         # these attributes are populated by _traverse_graph
         self.in_funcs = set()
         self.split_parents = []
+        self.split_branches = []
         self.matching_join = None
         # these attributes are populated by _postprocess
         self.is_inside_foreach = False
@@ -71,8 +86,57 @@ class DAGNode(object):
     def _expr_str(self, expr):
         return "%s.%s" % (expr.value.id, expr.attr)
 
-    def _parse(self, func_ast):
+    def _parse_switch_dict(self, dict_node):
+        switch_cases = {}
+
+        if isinstance(dict_node, ast.Dict):
+            for key, value in zip(dict_node.keys, dict_node.values):
+                case_key = None
+
+                # handle string literals
+                if hasattr(ast, "Str") and isinstance(key, ast.Str):
+                    case_key = key.s
+                elif isinstance(key, ast.Constant):
+                    case_key = key.value
+                elif isinstance(key, ast.Attribute):
+                    if isinstance(key.value, ast.Attribute) and isinstance(
+                        key.value.value, ast.Name
+                    ):
+                        # This handles self.config.some_key
+                        if key.value.value.id == "self":
+                            config_var = key.value.attr
+                            config_key = key.attr
+                            case_key = f"config:{config_var}.{config_key}"
+                        else:
+                            return None
+                    else:
+                        return None
+
+                # handle variables or other dynamic expressions - not allowed
+                elif isinstance(key, ast.Name):
+                    return None
+                else:
+                    # can't statically analyze this key
+                    return None
+
+                if case_key is None:
+                    return None
+
+                # extract the step name from the value
+                if isinstance(value, ast.Attribute) and isinstance(
+                    value.value, ast.Name
+                ):
+                    if value.value.id == "self":
+                        step_name = value.attr
+                        switch_cases[case_key] = step_name
+                    else:
+                        return None
+                else:
+                    return None
+
+        return switch_cases if switch_cases else None
 
+    def _parse(self, func_ast, lineno):
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
 
@@ -92,8 +156,39 @@ class DAGNode(object):
 
             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]
+            self.tail_next_lineno = lineno + tail.lineno - 1
+
+            # Check if first argument is a dictionary (switch case)
+            if (
+                len(tail.value.args) == 1
+                and isinstance(tail.value.args[0], ast.Dict)
+                and any(k.arg == "condition" for k in tail.value.keywords)
+            ):
+                # This is a switch statement
+                switch_cases = self._parse_switch_dict(tail.value.args[0])
+                condition_name = None
+
+                # Get condition parameter
+                for keyword in tail.value.keywords:
+                    if keyword.arg == "condition":
+                        if hasattr(ast, "Str") and isinstance(keyword.value, ast.Str):
+                            condition_name = keyword.value.s
+                        elif isinstance(keyword.value, ast.Constant) and isinstance(
+                            keyword.value.value, str
+                        ):
+                            condition_name = keyword.value.value
+                        break
+
+                if switch_cases and condition_name:
+                    self.type = "split-switch"
+                    self.condition = condition_name
+                    self.switch_cases = switch_cases
+                    self.out_funcs = list(switch_cases.values())
+                    self.invalid_tail_next = False
+                    return
+
+            else:
+                self.out_funcs = [e.attr for e in tail.value.args]
 
             keywords = dict(
                 (k.arg, getattr(k.value, "s", None)) for k in tail.value.keywords
@@ -129,10 +224,11 @@ class DAGNode(object):
             return
 
     def __str__(self):
-        return """*[{0.name} {0.type} (line {0.func_lineno})]*
+        return """*[{0.name} {0.type} ({0.source_file} line {0.func_lineno})]*
     in_funcs={in_funcs}
     out_funcs={out_funcs}
     split_parents={parents}
+    split_branches={branches}
     matching_join={matching_join}
     is_inside_foreach={is_inside_foreach}
     decorators={decos}
@@ -140,6 +236,7 @@ class DAGNode(object):
     has_tail_next={0.has_tail_next} (line {0.tail_next_lineno})
     invalid_tail_next={0.invalid_tail_next}
     foreach_param={0.foreach_param}
+    condition={0.condition}
     parallel_step={0.parallel_step}
     parallel_foreach={0.parallel_foreach}
     -> {out}""".format(
@@ -149,39 +246,43 @@ class DAGNode(object):
             out_funcs=", ".join("[%s]" % x for x in self.out_funcs),
             in_funcs=", ".join("[%s]" % x for x in self.in_funcs),
             parents=", ".join("[%s]" % x for x in self.split_parents),
+            branches=", ".join("[%s]" % x for x in self.split_branches),
             decos=" | ".join(map(str, self.decorators)),
             out=", ".join("[%s]" % x for x in self.out_funcs),
         )
 
 
-class StepVisitor(ast.NodeVisitor):
-    def __init__(self, nodes, flow):
-        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__)
-
-
 class FlowGraph(object):
     def __init__(self, flow):
         self.name = flow.__name__
         self.nodes = self._create_nodes(flow)
         self.doc = deindent_docstring(flow.__doc__)
+        # nodes sorted in topological order.
+        self.sorted_nodes = []
         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
-        ]
         nodes = {}
-        StepVisitor(nodes, flow).visit(root)
+        for element in dir(flow):
+            func = getattr(flow, element)
+            if callable(func) and hasattr(func, "is_step"):
+                source_file = inspect.getsourcefile(func)
+                source_lines, lineno = inspect.getsourcelines(func)
+                # This also works for code (strips out leading whitspace based on
+                # first line)
+                source_code = deindent_docstring("".join(source_lines))
+                function_ast = ast.parse(source_code).body[0]
+                node = DAGNode(
+                    function_ast,
+                    func.decorators,
+                    func.wrappers,
+                    func.config_decorators,
+                    func.__doc__,
+                    source_file,
+                    lineno,
+                )
+                nodes[element] = node
         return nodes
 
     def _postprocess(self):
@@ -196,18 +297,32 @@ class FlowGraph(object):
                 node.is_inside_foreach = True
 
     def _traverse_graph(self):
-        def traverse(node, seen, split_parents):
+        def traverse(node, seen, split_parents, split_branches):
+            add_split_branch = False
+            try:
+                self.sorted_nodes.remove(node.name)
+            except ValueError:
+                pass
+            self.sorted_nodes.append(node.name)
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
+                node.split_branches = split_branches
+                add_split_branch = True
                 split_parents = split_parents + [node.name]
+            elif node.type == "split-switch":
+                node.split_parents = split_parents
+                node.split_branches = split_branches
             elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
                     self[split_parents[-1]].matching_join = node.name
                     node.split_parents = split_parents
+                    node.split_branches = split_branches[:-1]
                     split_parents = split_parents[:-1]
+                    split_branches = split_branches[:-1]
             else:
                 node.split_parents = split_parents
+                node.split_branches = split_branches
 
             for n in node.out_funcs:
                 # graph may contain loops - ignore them
@@ -216,10 +331,15 @@ class FlowGraph(object):
                     if n in self:
                         child = self[n]
                         child.in_funcs.add(node.name)
-                        traverse(child, seen + [n], split_parents)
+                        traverse(
+                            child,
+                            seen + [n],
+                            split_parents,
+                            split_branches + ([n] if add_split_branch else []),
+                        )
 
         if "start" in self:
-            traverse(self["start"], [], [])
+            traverse(self["start"], [], [], [])
 
         # fix the order of in_funcs
         for node in self.nodes.values():
@@ -235,22 +355,42 @@ class FlowGraph(object):
         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(self[n]) for n in self.sorted_nodes)
 
     def output_dot(self):
         def edge_specs():
             for node in self.nodes.values():
-                for edge in node.out_funcs:
-                    yield "%s -> %s;" % (node.name, edge)
+                if node.type == "split-switch":
+                    # Label edges for switch cases
+                    for case_value, step_name in node.switch_cases.items():
+                        yield (
+                            '{0} -> {1} [label="{2}" color="blue" fontcolor="blue"];'.format(
+                                node.name, step_name, case_value
+                            )
+                        )
+                else:
+                    for edge in node.out_funcs:
+                        yield "%s -> %s;" % (node.name, edge)
 
         def node_specs():
             for node in self.nodes.values():
-                nodetype = "join" if node.num_args > 1 else node.type
-                yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
-                    node, type=nodetype
-                )
+                if node.type == "split-switch":
+                    # Hexagon shape for switch nodes
+                    condition_label = (
+                        f"switch: {node.condition}" if node.condition else "switch"
+                    )
+                    yield (
+                        '"{0.name}" '
+                        '[ label = <<b>{0.name}</b><br/><font point-size="9">{condition}</font>> '
+                        '  fontname = "Helvetica" '
+                        '  shape = "hexagon" '
+                        '  style = "filled" fillcolor = "lightgreen" ];'
+                    ).format(node, condition=condition_label)
+                else:
+                    nodetype = "join" if node.num_args > 1 else node.type
+                    yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
+                        node, type=nodetype
+                    )
 
         return (
             "digraph {0.name} {{\n"
@@ -262,7 +402,6 @@ class FlowGraph(object):
         )
 
     def output_steps(self):
-
         steps_info = {}
         graph_structure = []
 
@@ -275,6 +414,8 @@ class FlowGraph(object):
                 if node.parallel_foreach:
                     return "split-parallel"
                 return "split-foreach"
+            elif node.type == "split-switch":
+                return "split-switch"
             return "unknown"  # Should never happen
 
         def node_to_dict(name, node):
@@ -282,15 +423,26 @@ class FlowGraph(object):
                 "name": name,
                 "type": node_to_type(node),
                 "line": node.func_lineno,
+                "source_file": node.source_file,
                 "doc": node.doc,
                 "decorators": [
                     {
                         "name": deco.name,
-                        "attributes": deco.attributes,
+                        "attributes": to_pod(deco.attributes),
                         "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
                     }
                     for deco in node.decorators
                     if not deco.name.startswith("_")
+                ]
+                + [
+                    {
+                        "name": deco.decorator_name,
+                        "attributes": {"_args": deco._args, **deco._kwargs},
+                        "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
+                    }
+                    for deco in chain(node.wrappers, node.config_decorators)
                 ],
                 "next": node.out_funcs,
             }
@@ -298,6 +450,9 @@ class FlowGraph(object):
                 d["foreach_artifact"] = node.foreach_param
             elif d["type"] == "split-parallel":
                 d["num_parallel"] = node.num_parallel
+            elif d["type"] == "split-switch":
+                d["condition"] = node.condition
+                d["switch_cases"] = node.switch_cases
             if node.matching_join:
                 d["matching_join"] = node.matching_join
             return d
@@ -312,8 +467,8 @@ class FlowGraph(object):
                 steps_info[cur_name] = node_dict
                 resulting_list.append(cur_name)
 
-                if cur_node.type not in ("start", "linear", "join"):
-                    # We need to look at the different branches for this
+                node_type = node_to_type(cur_node)
+                if node_type in ("split-static", "split-foreach"):
                     resulting_list.append(
                         [
                             populate_block(s, cur_node.matching_join)
@@ -321,8 +476,21 @@ class FlowGraph(object):
                         ]
                     )
                     cur_name = cur_node.matching_join
+                elif node_type == "split-switch":
+                    all_paths = [
+                        populate_block(s, end_name)
+                        for s in cur_node.out_funcs
+                        if s != cur_name
+                    ]
+                    resulting_list.append(all_paths)
+                    cur_name = end_name
                 else:
-                    cur_name = cur_node.out_funcs[0]
+                    # handles only linear, start, and join steps.
+                    if cur_node.out_funcs:
+                        cur_name = cur_node.out_funcs[0]
+                    else:
+                        # handles terminal nodes or when we jump to 'end_name'.
+                        break
             return resulting_list
 
         graph_structure = populate_block("start", "end")

metaflow/includefile.py

@@ -1,14 +1,16 @@
 from collections import namedtuple
 import gzip
 
+import importlib
 import io
 import json
 import os
 
 from hashlib import sha1
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 from metaflow._vendor import click
+from metaflow._vendor import yaml
 
 from .exception import MetaflowException
 from .parameters import (
@@ -17,6 +19,9 @@ from .parameters import (
     Parameter,
     ParameterContext,
 )
+
+from .plugins import DATACLIENTS
+from .user_configs.config_options import ConfigInput
 from .util import get_username
 
 import functools
@@ -47,16 +52,7 @@ _DelayedExecContext = namedtuple(
 
 
 # From here on out, this is the IncludeFile implementation.
-from metaflow.plugins.datatools import Local, S3
-from metaflow.plugins.azure.includefile_support import Azure
-from metaflow.plugins.gcp.includefile_support import GS
-
-DATACLIENTS = {
-    "local": Local,
-    "s3": S3,
-    "azure": Azure,
-    "gs": GS,
-}
+_dict_dataclients = {d.TYPE: d for d in DATACLIENTS}
 
 
 class IncludedFile(object):
@@ -142,6 +138,7 @@ class FilePathClass(click.ParamType):
                 parameter_name=param.name,
                 logger=ctx.obj.echo,
                 ds_type=ctx.obj.datastore_impl.TYPE,
+                configs=None,
             )
 
         if len(value) > 0 and (value.startswith("{") or value.startswith('"{')):
@@ -167,7 +164,7 @@ class FilePathClass(click.ParamType):
                 "IncludeFile using a direct reference to a file in cloud storage is no "
                 "longer supported. Contact the Metaflow team if you need this supported"
             )
-            # if DATACLIENTS.get(path[:prefix_pos]) is None:
+            # if _dict_dataclients.get(path[:prefix_pos]) is None:
             #     self.fail(
             #         "IncludeFile: no handler for external file of type '%s' "
             #         "(given path is '%s')" % (path[:prefix_pos], path)
@@ -187,7 +184,7 @@ class FilePathClass(click.ParamType):
                     pass
             except OSError:
                 self.fail("IncludeFile: could not open file '%s' for reading" % path)
-            handler = DATACLIENTS.get(ctx.ds_type)
+            handler = _dict_dataclients.get(ctx.ds_type)
             if handler is None:
                 self.fail(
                     "IncludeFile: no data-client for datastore of type '%s'"
@@ -213,7 +210,7 @@ class FilePathClass(click.ParamType):
                         ctx.path,
                         ctx.is_text,
                         ctx.encoding,
-                        DATACLIENTS[ctx.handler_type],
+                        _dict_dataclients[ctx.handler_type],
                         ctx.echo,
                     )
                 )
@@ -249,29 +246,71 @@ class IncludeFile(Parameter):
     default : Union[str, Callable[ParameterContext, str]]
         Default path to a local file. A function
         implies that the parameter corresponds to a *deploy-time parameter*.
-    is_text : bool, default True
+    is_text : bool, optional, default None
         Convert the file contents to a string using the provided `encoding`.
-        If False, the artifact is stored in `bytes`.
-    encoding : str, optional, default 'utf-8'
-        Use this encoding to decode the file contexts if `is_text=True`.
-    required : bool, default False
+        If False, the artifact is stored in `bytes`. A value of None is equivalent to
+        True.
+    encoding : str, optional, default None
+        Use this encoding to decode the file contexts if `is_text=True`. A value of None
+        is equivalent to "utf-8".
+    required : bool, optional, default None
         Require that the user specified a value for the parameter.
-        `required=True` implies that the `default` is not used.
+        `required=True` implies that the `default` is not used. A value of None is
+        equivalent to False
     help : str, optional
         Help text to show in `run --help`.
     show_default : bool, default True
-        If True, show the default value in the help text.
+        If True, show the default value in the help text. A value of None is equivalent
+        to True.
+    parser : Union[str, Callable[[str], Any]], optional, default None
+        If a callable, it is a function that can parse the file contents
+        into any desired format. If a string, the string should refer to
+        a function (like "my_parser_package.my_parser.my_parser_function") which should
+        be able to parse the file contents. If the name starts with a ".", it is assumed
+        to be relative to "metaflow".
     """
 
     def __init__(
         self,
         name: str,
-        required: bool = False,
-        is_text: bool = True,
-        encoding: str = "utf-8",
+        required: Optional[bool] = None,
+        is_text: Optional[bool] = None,
+        encoding: Optional[str] = None,
         help: Optional[str] = None,
+        parser: Optional[Union[str, Callable[[str], Any]]] = None,
         **kwargs: Dict[str, str]
     ):
+        self._includefile_overrides = {}
+        if is_text is not None:
+            self._includefile_overrides["is_text"] = is_text
+        if encoding is not None:
+            self._includefile_overrides["encoding"] = encoding
+        self._parser = parser
+        # NOTA: Right now, there is an issue where these can't be overridden by config
+        # in all circumstances. Ignoring for now.
+        super(IncludeFile, self).__init__(
+            name,
+            required=required,
+            help=help,
+            type=FilePathClass(
+                self._includefile_overrides.get("is_text", True),
+                self._includefile_overrides.get("encoding", "utf-8"),
+            ),
+            **kwargs,
+        )
+
+    def init(self, ignore_errors=False):
+        super(IncludeFile, self).init(ignore_errors)
+
+        # This will use the values set explicitly in the args if present, else will
+        # use and remove from kwargs else will use True/utf-8
+        is_text = self._includefile_overrides.get(
+            "is_text", self.kwargs.pop("is_text", True)
+        )
+        encoding = self._includefile_overrides.get(
+            "encoding", self.kwargs.pop("encoding", "utf-8")
+        )
+
         # If a default is specified, it needs to be uploaded when the flow is deployed
         # (for example when doing a `step-functions create`) so we make the default
         # be a DeployTimeField. This means that it will be evaluated in two cases:
@@ -281,7 +320,7 @@ class IncludeFile(Parameter):
         # In the first case, we will need to fully upload the file whereas in the
         # second case, we can just return the string as the FilePath.convert method
         # will take care of evaluating things.
-        v = kwargs.get("default")
+        v = self.kwargs.get("default")
         if v is not None:
             # If the default is a callable, we have two DeployTimeField:
             #  - the callable nature of the default will require us to "call" the default
@@ -294,27 +333,32 @@ class IncludeFile(Parameter):
             # (call the default)
             if callable(v) and not isinstance(v, DeployTimeField):
                 # If default is a callable, make it a DeployTimeField (the inner one)
-                v = DeployTimeField(name, str, "default", v, return_str=True)
-            kwargs["default"] = DeployTimeField(
-                name,
+                v = DeployTimeField(self.name, str, "default", v, return_str=True)
+            self.kwargs["default"] = DeployTimeField(
+                self.name,
                 str,
                 "default",
                 IncludeFile._eval_default(is_text, encoding, v),
                 print_representation=v,
             )
 
-        super(IncludeFile, self).__init__(
-            name,
-            required=required,
-            help=help,
-            type=FilePathClass(is_text, encoding),
-            **kwargs,
-        )
-
     def load_parameter(self, v):
         if v is None:
             return v
-        return v.decode(self.name, var_type="Parameter")
+
+        # Get the raw content from the file
+        content = v.decode(self.name, var_type="Parameter")
+        # If a parser is specified, use it to parse the content
+        if self._parser is not None:
+            try:
+                return ConfigInput._call_parser(self._parser, content)
+            except Exception as e:
+                raise MetaflowException(
+                    "Failed to parse content in parameter '%s' using parser: %s"
+                    % (self.name, str(e))
+                ) from e
+
+        return content
 
     @staticmethod
     def _eval_default(is_text, encoding, default_path):
@@ -425,7 +469,7 @@ class UploaderV1:
         if prefix_pos < 0:
             raise MetaflowException("Malformed URL: '%s'" % url)
         prefix = url[:prefix_pos]
-        handler = DATACLIENTS.get(prefix)
+        handler = _dict_dataclients.get(prefix)
         if handler is None:
             raise MetaflowException("Could not find data client for '%s'" % prefix)
         return handler
@@ -437,7 +481,7 @@ class UploaderV2:
     @classmethod
     def encode_url(cls, url_type, url, **kwargs):
         return_value = {
-            "note": "Internal representation of IncludeFile(%s)" % url,
+            "note": "Internal representation of IncludeFile",
             "type": cls.file_type,
             "sub-type": url_type,
             "url": url,
@@ -450,7 +494,7 @@ class UploaderV2:
         r = UploaderV1.store(flow_name, path, is_text, encoding, handler, echo)
 
         # In V2, we store size for faster access
-        r["note"] = "Internal representation of IncludeFile(%s)" % path
+        r["note"] = "Internal representation of IncludeFile"
         r["type"] = cls.file_type
         r["sub-type"] = "uploaded"
         r["size"] = os.stat(path).st_size