Decorators now declare the configuration parameters they need This allows for inspection/visibility, and will be essential in checkpointing/storage options. We want decorators to declare these upfront, rather than relying on the internals of the decorators to break, enabling clearer error messages for more usable configuration-driven pipelines. Furthermore, we also allow for some "optional" configuration parameters. These enable us to provide a default. The use-case here is config.when -- which currently defaults to None. Allowing this from the start is not ideal, but it's what we did and users rely on it. There *is* an out, however. If someone is using the base `@config`, they were able to declare a lambda function, which yields no insight into the parameters required by the function. To bypass this, we pass None out, and the framework knows not to filter the required parameters.
diff --git a/hamilton/function_modifiers/base.py b/hamilton/function_modifiers/base.py index 9cdcfc3..cd89604 100644 --- a/hamilton/function_modifiers/base.py +++ b/hamilton/function_modifiers/base.py
@@ -2,7 +2,7 @@ import collections import functools import logging -from typing import Any, Callable, Collection, Dict, List, Tuple +from typing import Any, Callable, Collection, Dict, List, Optional, Tuple from hamilton import node, registry @@ -112,16 +112,45 @@ setattr(fn, lifecycle_name, [self]) return fn + def required_config(self) -> Optional[List[str]]: + """Declares the required configuration keys for this decorator. + Note that these configuration keys will be filtered and passed to the `configuration` + parameter of the functions that this decorator uses. + + Note that this currently allows for a "escape hatch". + That is, returning None from this function. + + :return: A list of the required configuration keys. + """ + return [] + + def optional_config(self) -> Dict[str, Any]: + """Declares the optional configuration keys for this decorator. + These are configuration keys that can be used by the decorator, but are not required. + Along with these we have *defaults*, which we will use to pass to the config. + + :return: + """ + return {} + + @property + def name(self) -> str: + """Name of the decorator. + + :return: The name of the decorator + """ + return self.__class__.__name__ + class NodeResolver(NodeTransformLifecycle): """Decorator to resolve a nodes function. Can modify anything about the function and is run at DAG creation time.""" @abc.abstractmethod - def resolve(self, fn: Callable, configuration: Dict[str, Any]) -> Callable: + def resolve(self, fn: Callable, config: Dict[str, Any]) -> Callable: """Determines what a function resolves to. Returns None if it should not be included in the DAG. :param fn: Function to resolve - :param configuration: DAG config + :param config: DAG config :return: A name if it should resolve to something. Otherwise None. """ pass @@ -339,7 +368,7 @@ class DefaultNodeResolver(NodeResolver): - def resolve(self, fn: Callable, configuration: Dict[str, Any]) -> Callable: + def resolve(self, fn: Callable, config: Dict[str, Any]) -> Callable: return fn def validate(self, fn): @@ -361,6 +390,43 @@ return node_ +def merge_configs( + name_for_error: str, + config: Dict[str, Any], + config_required: List[str], + config_optional_with_defaults: Dict[str, Any], +) -> Dict[str, Any]: + missing_keys = ( + set(config_required) - set(config.keys()) - set(config_optional_with_defaults.keys()) + ) + if len(missing_keys) > 0: + raise ValueError( + f"The following configurations are required by {name_for_error}: {missing_keys} b" + ) + config_out = {key: config[key] for key in config_required} + for key in config_optional_with_defaults: + config_out[key] = config.get(key, config_optional_with_defaults[key]) + return config_out + + +def filter_config(config: Dict[str, Any], decorator: NodeTransformLifecycle) -> Dict[str, Any]: + """Filters the config to only include the keys in config_required + TODO -- break this into two so we can make it easier to test. + + :param config: The config to filter + :param config_required: The keys to include + :param decorator: The decorator that is utilizing the configuration + :return: The filtered config + """ + config_required = decorator.required_config() + config_optional_with_defaults = decorator.optional_config() + if config_required is None: + # This is an out to allow for backwards compatibility for the config.resolve decorator + # Note this is an internal API, but we made the config with the `resolve` parameter public + return config + return merge_configs(decorator.name, config, config_required, config_optional_with_defaults) + + def resolve_nodes(fn: Callable, config: Dict[str, Any]) -> Collection[node.Node]: """Gets a list of nodes from a function. This is meant to be an abstraction between the node and the function that it implements. This will end up coordinating with the decorators we build @@ -388,24 +454,26 @@ 5. Return the final list of nodes. :param fn: Function to input. + :param config: Configuratino to use -- this can be used by decorators to specify + which configuration they need. :return: A list of nodes into which this function transforms. """ node_resolvers = getattr(fn, NodeResolver.get_lifecycle_name(), [DefaultNodeResolver()]) for resolver in node_resolvers: - fn = resolver.resolve(fn, config) + fn = resolver.resolve(fn, config=filter_config(config, resolver)) if fn is None: return [] (node_creator,) = getattr(fn, NodeCreator.get_lifecycle_name(), [DefaultNodeCreator()]) - nodes = node_creator.generate_nodes(fn, config) + nodes = node_creator.generate_nodes(fn, filter_config(config, node_creator)) if hasattr(fn, NodeExpander.get_lifecycle_name()): (node_expander,) = getattr(fn, NodeExpander.get_lifecycle_name(), [DefaultNodeExpander()]) - nodes = node_expander.transform_dag(nodes, config, fn) + nodes = node_expander.transform_dag(nodes, filter_config(config, node_expander), fn) node_transformers = getattr(fn, NodeTransformer.get_lifecycle_name(), []) for dag_modifier in node_transformers: - nodes = dag_modifier.transform_dag(nodes, config, fn) + nodes = dag_modifier.transform_dag(nodes, filter_config(config, dag_modifier), fn) node_decorators = getattr(fn, NodeDecorator.get_lifecycle_name(), [DefaultNodeDecorator()]) for node_decorator in node_decorators: - nodes = node_decorator.transform_dag(nodes, config, fn) + nodes = node_decorator.transform_dag(nodes, filter_config(config, node_decorator), fn) return nodes
diff --git a/hamilton/function_modifiers/configuration.py b/hamilton/function_modifiers/configuration.py index 75532b7..e49ad20 100644 --- a/hamilton/function_modifiers/configuration.py +++ b/hamilton/function_modifiers/configuration.py
@@ -1,4 +1,4 @@ -from typing import Any, Callable, Collection, Dict +from typing import Any, Callable, Collection, Dict, List, Optional from . import base @@ -13,17 +13,31 @@ That said, you can have functions that *only* exist in certain configurations without worrying about it. """ - def __init__(self, resolves: Callable[[Dict[str, Any]], bool], target_name: str = None): + def __init__( + self, + resolves: Callable[[Dict[str, Any]], bool], + target_name: str = None, + config_used: List[str] = None, + ): self.does_resolve = resolves self.target_name = target_name + self._config_used = config_used + + def required_config(self) -> Optional[List[str]]: + """Nothing is currently required""" + return [] # All of these can default to None + + def optional_config(self) -> Dict[str, Any]: + """Everything is optional with None as the required value""" + return {key: None for key in self._config_used} def _get_function_name(self, fn: Callable) -> str: if self.target_name is not None: return self.target_name return base.sanitize_function_name(fn.__name__) - def resolve(self, fn, configuration: Dict[str, Any]) -> Callable: - if not self.does_resolve(configuration): + def resolve(self, fn, config: Dict[str, Any]) -> Callable: + if not self.does_resolve(config): return None fn.__name__ = self._get_function_name(fn) # TODO -- copy function to not mutate it return fn @@ -45,7 +59,7 @@ def resolves(configuration: Dict[str, Any]) -> bool: return all(value == configuration.get(key) for key, value in key_value_pairs.items()) - return config(resolves, target_name=name) + return config(resolves, target_name=name, config_used=list(key_value_pairs.keys())) @staticmethod def when_not(name=None, **key_value_pairs: Any) -> "config": @@ -58,11 +72,12 @@ def resolves(configuration: Dict[str, Any]) -> bool: return all(value != configuration.get(key) for key, value in key_value_pairs.items()) - return config(resolves, target_name=name) + return config(resolves, target_name=name, config_used=list(key_value_pairs.keys())) @staticmethod def when_in(name=None, **key_value_group_pairs: Collection[Any]) -> "config": - """Yields a decorator that resolves the function if all of the keys are equal to one of items in the list of values. + """Yields a decorator that resolves the function if all of the + values corresponding to the config keys are equal to one of items in the list of values. :param key_value_group_pairs: pairs of key-value mappings where the value is a list of possible values :return: a configuration decorator @@ -73,7 +88,7 @@ configuration.get(key) in value for key, value in key_value_group_pairs.items() ) - return config(resolves, target_name=name) + return config(resolves, target_name=name, config_used=list(key_value_group_pairs.keys())) @staticmethod def when_not_in(**key_value_group_pairs: Collection[Any]) -> "config": @@ -100,4 +115,4 @@ configuration.get(key) not in value for key, value in key_value_group_pairs.items() ) - return config(resolves) + return config(resolves, config_used=list(key_value_group_pairs.keys()))
diff --git a/hamilton/function_modifiers/macros.py b/hamilton/function_modifiers/macros.py index a133c56..70b1e01 100644 --- a/hamilton/function_modifiers/macros.py +++ b/hamilton/function_modifiers/macros.py
@@ -226,6 +226,13 @@ ) ] + def require_config(self) -> List[str]: + """Returns the configuration parameters that this model requires + + :return: Just the one config param used by this model + """ + return [self.config_param] + class model(dynamic_transform): """Model, same as a dynamic transform"""