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

metaflow/user_configs/config_parameters.py

@@ -1,9 +1,22 @@
-import collections.abc
+import inspect
 import json
+import collections.abc
+import copy
 import os
 import re
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 
 from ..exception import MetaflowException
@@ -36,25 +49,21 @@ if TYPE_CHECKING:
 
 #     return tracefunc_closure
 
-CONFIG_FILE = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "CONFIG_PARAMETERS"
-)
-
 ID_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
 
 UNPACK_KEY = "_unpacked_delayed_"
 
 
 def dump_config_values(flow: "FlowSpec"):
-    from ..flowspec import _FlowState  # Prevent circular import
+    from ..flowspec import FlowStateItems  # Prevent circular import
 
-    configs = flow._flow_state.get(_FlowState.CONFIGS)
+    configs = flow._flow_state[FlowStateItems.CONFIGS]
     if configs:
         return {"user_configs": configs}
     return {}
 
 
-class ConfigValue(collections.abc.Mapping):
+class ConfigValue(collections.abc.Mapping, dict):
     """
     ConfigValue is a thin wrapper around an arbitrarily nested dictionary-like
     configuration object. It allows you to access elements of this nested structure
@@ -69,8 +78,67 @@ class ConfigValue(collections.abc.Mapping):
     # Thin wrapper to allow configuration values to be accessed using a "." notation
     # as well as a [] notation.
 
-    def __init__(self, data: Dict[str, Any]):
-        self._data = data
+    # We inherit from dict to allow the isinstanceof check to work easily and also
+    # to provide a simple json dumps functionality.
+
+    def __init__(self, data: Union["ConfigValue", Dict[str, Any]]):
+        self._data = {k: self._construct(v) for k, v in data.items()}
+
+        # Enable json dumps
+        dict.__init__(self, self._data)
+
+    @classmethod
+    def fromkeys(cls, iterable: Iterable, value: Any = None) -> "ConfigValue":
+        """
+        Creates a new ConfigValue object from the given iterable and value.
+
+        Parameters
+        ----------
+        iterable : Iterable
+            Iterable to create the ConfigValue from.
+        value : Any, optional
+            Value to set for each key in the iterable.
+
+        Returns
+        -------
+        ConfigValue
+            A new ConfigValue object.
+        """
+        return cls(dict.fromkeys(iterable, value))
+
+    def to_dict(self) -> Dict[Any, Any]:
+        """
+        Returns a dictionary representation of this configuration object.
+
+        Returns
+        -------
+        Dict[Any, Any]
+            Dictionary equivalent of this configuration object.
+        """
+        return self._to_dict(self._data)
+
+    def copy(self) -> "ConfigValue":
+        return self.__copy__()
+
+    def clear(self) -> None:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def update(self, *args, **kwargs) -> None:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def setdefault(self, key: Any, default: Any = None) -> Any:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def pop(self, key: Any, default: Any = None) -> Any:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def popitem(self) -> Tuple[Any, Any]:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
 
     def __getattr__(self, key: str) -> Any:
         """
@@ -115,33 +183,96 @@ class ConfigValue(collections.abc.Mapping):
         Any
             Element of the configuration
         """
-        value = self._data[key]
-        if isinstance(value, dict):
-            value = ConfigValue(value)
-        return value
+        return self._data[key]
 
-    def __len__(self):
+    def __setitem__(self, key: Any, value: Any) -> None:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def __delattr__(self, key) -> None:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def __delitem__(self, key: Any) -> None:
+        # Prevent configuration modification
+        raise TypeError("ConfigValue is immutable")
+
+    def __len__(self) -> int:
         return len(self._data)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator:
         return iter(self._data)
 
-    def __repr__(self):
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, ConfigValue):
+            return self._data == other._data
+        if isinstance(other, dict):
+            return self._data == other
+        return False
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
+
+    def __copy__(self) -> "ConfigValue":
+        cls = self.__class__
+        result = cls.__new__(cls)
+        result.__dict__.update({k: copy.copy(v) for k, v in self.__dict__.items()})
+        return result
+
+    def __repr__(self) -> str:
         return repr(self._data)
 
-    def __str__(self):
-        return json.dumps(self._data)
+    def __str__(self) -> str:
+        return str(self._data)
 
-    def to_dict(self) -> Dict[Any, Any]:
+    def __dir__(self) -> Iterable[str]:
+        return dir(type(self)) + [k for k in self._data.keys() if ID_PATTERN.match(k)]
+
+    def __contains__(self, key: Any) -> bool:
+        try:
+            self[key]
+        except KeyError:
+            return False
+        return True
+
+    def keys(self):
         """
-        Returns a dictionary representation of this configuration object.
+        Returns the keys of this configuration object.
 
         Returns
         -------
-        Dict[Any, Any]
-            Dictionary equivalent of this configuration object.
+        Any
+            Keys of this configuration object.
         """
-        return dict(self._data)
+        return self._data.keys()
+
+    @classmethod
+    def _construct(cls, obj: Any) -> Any:
+        # Internal method to construct a ConfigValue so that all mappings internally
+        # are also converted to ConfigValue
+        if isinstance(obj, ConfigValue):
+            v = obj
+        elif isinstance(obj, collections.abc.Mapping):
+            v = ConfigValue({k: cls._construct(v) for k, v in obj.items()})
+        elif isinstance(obj, (list, tuple, set)):
+            v = type(obj)([cls._construct(x) for x in obj])
+        else:
+            v = obj
+        return v
+
+    @classmethod
+    def _to_dict(cls, obj: Any) -> Any:
+        # Internal method to convert all nested mappings to dicts
+        if isinstance(obj, collections.abc.Mapping):
+            v = {k: cls._to_dict(v) for k, v in obj.items()}
+        elif isinstance(obj, (list, tuple, set)):
+            v = type(obj)([cls._to_dict(x) for x in obj])
+        else:
+            v = obj
+        return v
+
+    def __reduce__(self):
+        return (self.__class__, (self.to_dict(),))
 
 
 class DelayEvaluator(collections.abc.Mapping):
@@ -157,8 +288,9 @@ class DelayEvaluator(collections.abc.Mapping):
         of _unpacked_delayed_*
     """
 
-    def __init__(self, ex: str):
+    def __init__(self, ex: str, saved_globals: Optional[Dict[str, Any]] = None):
         self._config_expr = ex
+        self._globals = saved_globals
         if ID_PATTERN.match(self._config_expr):
             # This is a variable only so allow things like config_expr("config").var
             self._is_var_only = True
@@ -166,6 +298,21 @@ class DelayEvaluator(collections.abc.Mapping):
         else:
             self._is_var_only = False
             self._access = None
+        self._cached_expr = None
+
+    def __copy__(self):
+        c = DelayEvaluator(self._config_expr)
+        c._access = self._access.copy() if self._access is not None else None
+        # Globals are not copied -- always kept as a reference
+        return c
+
+    def __deepcopy__(self, memo):
+        c = DelayEvaluator(self._config_expr)
+        c._access = (
+            copy.deepcopy(self._access, memo) if self._access is not None else None
+        )
+        # Globals are not copied -- always kept as a reference
+        return c
 
     def __iter__(self):
         yield "%s%d" % (UNPACK_KEY, id(self))
@@ -175,8 +322,15 @@ class DelayEvaluator(collections.abc.Mapping):
             return self
         if self._access is None:
             raise KeyError(key)
-        self._access.append(key)
-        return self
+
+        # Make a copy so that we can support something like
+        # foo = delay_evaluator["blah"]
+        # bar = delay_evaluator["baz"]
+        # and don't end up with a access list that contains both "blah" and "baz"
+        c = self.__copy__()
+        c._access.append(key)
+        c._cached_expr = None
+        return c
 
     def __len__(self):
         return 1
@@ -184,11 +338,13 @@ class DelayEvaluator(collections.abc.Mapping):
     def __getattr__(self, name):
         if self._access is None:
             raise AttributeError(name)
-        self._access.append(name)
-        return self
+        c = self.__copy__()
+        c._access.append(name)
+        c._cached_expr = None
+        return c
 
     def __call__(self, ctx=None, deploy_time=False):
-        from ..flowspec import _FlowState  # Prevent circular import
+        from ..flowspec import FlowStateItems  # Prevent circular import
 
         # Two additional arguments are only used by DeployTimeField which will call
         # this function with those two additional arguments. They are ignored.
@@ -199,7 +355,9 @@ class DelayEvaluator(collections.abc.Mapping):
                 "Config object can only be used directly in the FlowSpec defining them "
                 "(or their flow decorators)."
             )
-        if self._access is not None:
+        if self._cached_expr is not None:
+            to_eval_expr = self._cached_expr
+        elif self._access is not None:
             # Build the final expression by adding all the fields in access as . fields
             access_list = [self._config_expr]
             for a in self._access:
@@ -212,27 +370,24 @@ class DelayEvaluator(collections.abc.Mapping):
                     raise MetaflowException(
                         "Field '%s' of type '%s' is not supported" % (str(a), type(a))
                     )
-            self._config_expr = ".".join(access_list)
+            to_eval_expr = self._cached_expr = ".".join(access_list)
+        else:
+            to_eval_expr = self._cached_expr = self._config_expr
         # Evaluate the expression setting the config values as local variables
         try:
             return eval(
-                self._config_expr,
-                globals(),
+                to_eval_expr,
+                self._globals or globals(),
                 {
-                    k: ConfigValue(v)
-                    for k, v in flow_cls._flow_state.get(_FlowState.CONFIGS, {}).items()
+                    k: ConfigValue(v) if v is not None else None
+                    for k, v in flow_cls._flow_state[FlowStateItems.CONFIGS].items()
                 },
             )
         except NameError as e:
-            potential_config_name = self._config_expr.split(".")[0]
-            if potential_config_name not in flow_cls._flow_state.get(
-                _FlowState.CONFIGS, {}
-            ):
-                raise MetaflowException(
-                    "Config '%s' not found in the flow (maybe not required and not "
-                    "provided?)" % potential_config_name
-                ) from e
-            raise
+            raise MetaflowException(
+                "Config expression '%s' could not be evaluated: %s"
+                % (to_eval_expr, str(e))
+            ) from e
 
 
 def config_expr(expr: str) -> DelayEvaluator:
@@ -262,7 +417,10 @@ def config_expr(expr: str) -> DelayEvaluator:
     expr : str
         Expression using the config values.
     """
-    return DelayEvaluator(expr)
+    # Get globals where the expression is defined so that the user can use
+    # something like `config_expr("my_func()")` in the expression.
+    parent_globals = inspect.currentframe().f_back.f_globals
+    return DelayEvaluator(expr, saved_globals=parent_globals)
 
 
 class Config(Parameter, collections.abc.Mapping):
@@ -329,7 +487,6 @@ class Config(Parameter, collections.abc.Mapping):
         parser: Optional[Union[str, Callable[[str], Dict[Any, Any]]]] = None,
         **kwargs: Dict[str, str]
     ):
-
         if default is not None and default_value is not None:
             raise MetaflowException(
                 "For config '%s', you can only specify default or default_value, not both"
@@ -350,9 +507,7 @@ class Config(Parameter, collections.abc.Mapping):
         self._delayed_evaluator = None
 
     def load_parameter(self, v):
-        if v is None:
-            return None
-        return ConfigValue(v)
+        return ConfigValue(v) if v is not None else None
 
     def _store_value(self, v: Any) -> None:
         self._computed_value = v
@@ -385,23 +540,30 @@ class Config(Parameter, collections.abc.Mapping):
         return DelayEvaluator(self.name.lower())[key]
 
 
-def resolve_delayed_evaluator(v: Any, ignore_errors: bool = False) -> Any:
+def resolve_delayed_evaluator(
+    v: Any, ignore_errors: bool = False, to_dict: bool = False
+) -> Any:
     # NOTE: We don't ignore errors in downstream calls because we want to have either
     # all or nothing for the top-level call by the user.
     try:
         if isinstance(v, DelayEvaluator):
-            return v()
+            to_return = v()
+            if to_dict and isinstance(to_return, ConfigValue):
+                to_return = to_return.to_dict()
+            return to_return
         if isinstance(v, dict):
             return {
-                resolve_delayed_evaluator(k): resolve_delayed_evaluator(v)
+                resolve_delayed_evaluator(
+                    k, to_dict=to_dict
+                ): resolve_delayed_evaluator(v, to_dict=to_dict)
                 for k, v in v.items()
             }
         if isinstance(v, list):
-            return [resolve_delayed_evaluator(x) for x in v]
+            return [resolve_delayed_evaluator(x, to_dict=to_dict) for x in v]
         if isinstance(v, tuple):
-            return tuple(resolve_delayed_evaluator(x) for x in v)
+            return tuple(resolve_delayed_evaluator(x, to_dict=to_dict) for x in v)
         if isinstance(v, set):
-            return {resolve_delayed_evaluator(x) for x in v}
+            return {resolve_delayed_evaluator(x, to_dict=to_dict) for x in v}
         return v
     except Exception as e:
         if ignore_errors:
@@ -426,7 +588,7 @@ def unpack_delayed_evaluator(
         else:
             # k.startswith(UNPACK_KEY)
             try:
-                new_vals = resolve_delayed_evaluator(v)
+                new_vals = resolve_delayed_evaluator(v, to_dict=True)
                 new_keys.extend(new_vals.keys())
                 result.update(new_vals)
             except Exception as e:

metaflow/user_decorators/common.py

@@ -0,0 +1,144 @@
+from typing import Dict, Optional, List, Tuple
+
+
+class _TrieNode:
+    def __init__(
+        self, parent: Optional["_TrieNode"] = None, component: Optional[str] = None
+    ):
+        self.parent = parent
+        self.component = component
+        self.children = {}  # type: Dict[str, "_TrieNode"]
+        self.total_children = 0
+        self.value = None
+        self.end_value = None
+
+    def traverse(self, value: type) -> Optional["_TrieNode"]:
+        if self.total_children == 0:
+            self.end_value = value
+        else:
+            self.end_value = None
+        self.total_children += 1
+
+    def remove_child(self, child_name: str) -> bool:
+        if child_name in self.children:
+            del self.children[child_name]
+            self.total_children -= 1
+            return True
+        return False
+
+
+class ClassPath_Trie:
+    def __init__(self):
+        self.root = _TrieNode(None, None)
+        self.inited = False
+        self._value_to_node = {}  # type: Dict[type, _TrieNode]
+
+    def init(self, initial_nodes: Optional[List[Tuple[str, type]]] = None):
+        # We need to do this so we can delay import of STEP_DECORATORS
+        self.inited = True
+        for classpath_name, value in initial_nodes or []:
+            self.insert(classpath_name, value)
+
+    def insert(self, classpath_name: str, value: type):
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            node = node.children.setdefault(c, _TrieNode(node, c))
+            node.traverse(value)
+        node.total_children -= (
+            1  # We do not count the last node as having itself as a child
+        )
+        node.value = value
+        self._value_to_node[value] = node
+
+    def search(self, classpath_name: str) -> Optional[type]:
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            if c not in node.children:
+                return None
+            node = node.children[c]
+        return node.value
+
+    def remove(self, classpath_name: str):
+        components = list(reversed(classpath_name.split(".")))
+
+        def _remove(node: _TrieNode, components, depth):
+            if depth == len(components):
+                if node.value is not None:
+                    del self._value_to_node[node.value]
+                    node.value = None
+                    return len(node.children) == 0
+                return False
+            c = components[depth]
+            if c not in node.children:
+                return False
+            did_delete_child = _remove(node.children[c], components, depth + 1)
+            if did_delete_child:
+                node.remove_child(c)
+                if node.total_children == 1:
+                    # If we have one total child left, we have at least one
+                    # child and that one has an end_value
+                    for child in node.children.values():
+                        assert (
+                            child.end_value
+                        ), "Node with one child must have an end_value"
+                        node.end_value = child.end_value
+                return node.total_children == 0
+            return False
+
+        _remove(self.root, components, 0)
+
+    def unique_prefix_value(self, classpath_name: str) -> Optional[type]:
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            if c not in node.children:
+                return None
+            node = node.children[c]
+        # If we reach here, it means the classpath_name is a prefix.
+        # We check if it has only one path forward (end_value will be non None)
+        # If value is not None, we also consider this to be a unique "prefix"
+        # This happens since this trie is also filled with metaflow default decorators
+        return node.end_value or node.value
+
+    def unique_prefix_for_type(self, value: type) -> Optional[str]:
+        node = self._value_to_node.get(value, None)
+        if node is None:
+            return None
+        components = []
+        while node:
+            if node.end_value == value:
+                components = []
+            if node.component is not None:
+                components.append(node.component)
+            node = node.parent
+        return ".".join(components)
+
+    def get_unique_prefixes(self) -> Dict[str, type]:
+        """
+        Get all unique prefixes in the trie.
+
+        Returns
+        -------
+        List[str]
+            A list of unique prefixes.
+        """
+        to_return = {}
+
+        def _collect(node, current_prefix):
+            if node.end_value is not None:
+                to_return[current_prefix] = node.end_value
+                # We stop there and don't look further since we found the unique prefix
+                return
+            if node.value is not None:
+                to_return[current_prefix] = node.value
+                # We continue to look for more unique prefixes
+            for child_name, child_node in node.children.items():
+                _collect(
+                    child_node,
+                    f"{current_prefix}.{child_name}" if current_prefix else child_name,
+                )
+
+        _collect(self.root, "")
+        return {".".join(reversed(k.split("."))): v for k, v in to_return.items()}