Add superseding flag logic (#1078)

This PR adds a superseding relationship between flags, specified by the `_SUPERSEDING_FLAGS` dictionary in `_global_config.py` --------- Signed-off-by: Elliot Gunton <elliotgunton@gmail.com>
argoproj-labs · May 31, 2024 · fd79fbf · fd79fbf
1 parent fe4d5c9
commit fd79fbf
Show file tree

Hide file tree

Showing 16 changed files with 54 additions and 79 deletions.
diff --git a/docs/examples/workflows/experimental/new_dag_decorator_artifacts.md b/docs/examples/workflows/experimental/new_dag_decorator_artifacts.md
@@ -13,8 +13,6 @@
     from hera.shared import global_config
     from hera.workflows import Artifact, Input, Output, Workflow
 
-    global_config.experimental_features["script_annotations"] = True
-    global_config.experimental_features["script_pydantic_io"] = True
     global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/docs/examples/workflows/experimental/new_dag_decorator_inner_dag.md b/docs/examples/workflows/experimental/new_dag_decorator_inner_dag.md
@@ -11,8 +11,6 @@
     from hera.shared import global_config
     from hera.workflows import Input, Output, Workflow
 
-    global_config.experimental_features["script_annotations"] = True
-    global_config.experimental_features["script_pydantic_io"] = True
     global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/docs/examples/workflows/experimental/new_dag_decorator_params.md b/docs/examples/workflows/experimental/new_dag_decorator_params.md
@@ -14,8 +14,6 @@
     from hera.shared import global_config
     from hera.workflows import Input, Output, Parameter, Workflow
 
-    global_config.experimental_features["script_annotations"] = True
-    global_config.experimental_features["script_pydantic_io"] = True
     global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/docs/examples/workflows/experimental/new_decorators_basic_script.md b/docs/examples/workflows/experimental/new_decorators_basic_script.md
@@ -11,8 +11,6 @@
     from hera.shared import global_config
     from hera.workflows import Input, Output, WorkflowTemplate
 
-    global_config.experimental_features["script_annotations"] = True
-    global_config.experimental_features["script_pydantic_io"] = True
     global_config.experimental_features["decorator_syntax"] = True
 
     w = WorkflowTemplate(name="my-template")

diff --git a/docs/examples/workflows/experimental/new_steps_decorator_with_parallel_steps.md b/docs/examples/workflows/experimental/new_steps_decorator_with_parallel_steps.md
@@ -13,8 +13,6 @@
     from hera.shared import global_config
     from hera.workflows import Input, Output, Parameter, Workflow, parallel
 
-    global_config.experimental_features["script_annotations"] = True
-    global_config.experimental_features["script_pydantic_io"] = True
     global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/docs/examples/workflows/experimental/script_runner_io.md b/docs/examples/workflows/experimental/script_runner_io.md
@@ -23,7 +23,6 @@
     except ImportError:
         from typing_extensions import Annotated  # type: ignore
 
-    global_config.experimental_features["script_annotations"] = True
     global_config.experimental_features["script_pydantic_io"] = True
 
 

diff --git a/examples/workflows/experimental/new_dag_decorator_artifacts.py b/examples/workflows/experimental/new_dag_decorator_artifacts.py
@@ -3,8 +3,6 @@
 from hera.shared import global_config
 from hera.workflows import Artifact, Input, Output, Workflow
 
-global_config.experimental_features["script_annotations"] = True
-global_config.experimental_features["script_pydantic_io"] = True
 global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/examples/workflows/experimental/new_dag_decorator_inner_dag.py b/examples/workflows/experimental/new_dag_decorator_inner_dag.py
@@ -1,8 +1,6 @@
 from hera.shared import global_config
 from hera.workflows import Input, Output, Workflow
 
-global_config.experimental_features["script_annotations"] = True
-global_config.experimental_features["script_pydantic_io"] = True
 global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/examples/workflows/experimental/new_dag_decorator_params.py b/examples/workflows/experimental/new_dag_decorator_params.py
@@ -4,8 +4,6 @@
 from hera.shared import global_config
 from hera.workflows import Input, Output, Parameter, Workflow
 
-global_config.experimental_features["script_annotations"] = True
-global_config.experimental_features["script_pydantic_io"] = True
 global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/examples/workflows/experimental/new_decorators_basic_script.py b/examples/workflows/experimental/new_decorators_basic_script.py
@@ -1,8 +1,6 @@
 from hera.shared import global_config
 from hera.workflows import Input, Output, WorkflowTemplate
 
-global_config.experimental_features["script_annotations"] = True
-global_config.experimental_features["script_pydantic_io"] = True
 global_config.experimental_features["decorator_syntax"] = True
 
 w = WorkflowTemplate(name="my-template")

diff --git a/examples/workflows/experimental/new_steps_decorator_with_parallel_steps.py b/examples/workflows/experimental/new_steps_decorator_with_parallel_steps.py
@@ -3,8 +3,6 @@
 from hera.shared import global_config
 from hera.workflows import Input, Output, Parameter, Workflow, parallel
 
-global_config.experimental_features["script_annotations"] = True
-global_config.experimental_features["script_pydantic_io"] = True
 global_config.experimental_features["decorator_syntax"] = True
 
 

diff --git a/examples/workflows/experimental/script_runner_io.py b/examples/workflows/experimental/script_runner_io.py
@@ -13,7 +13,6 @@
 except ImportError:
     from typing_extensions import Annotated  # type: ignore
 
-global_config.experimental_features["script_annotations"] = True
 global_config.experimental_features["script_pydantic_io"] = True
 
 

diff --git a/src/hera/shared/_global_config.py b/src/hera/shared/_global_config.py
@@ -165,3 +165,23 @@ def _set_defaults(cls, values):
 
 GlobalConfig = global_config = _GlobalConfig()
 register_pre_build_hook = global_config.register_pre_build_hook
+
+_SCRIPT_ANNOTATIONS_FLAG = "script_annotations"
+_SCRIPT_PYDANTIC_IO_FLAG = "script_pydantic_io"
+_DECORATOR_SYNTAX_FLAG = "decorator_syntax"
+
+# A dictionary where each key is a flag that has a list of flags which supersede it, hence
+# the given flag key can also be switched on by any of the flags in the list. Using simple flat lists
+# for now, otherwise with many superseding flags we may want to have a recursive structure.
+_SUPERSEDING_FLAGS: Dict[str, List] = {
+    _SCRIPT_ANNOTATIONS_FLAG: [_SCRIPT_PYDANTIC_IO_FLAG, _DECORATOR_SYNTAX_FLAG],
+    _SCRIPT_PYDANTIC_IO_FLAG: [_DECORATOR_SYNTAX_FLAG],
+    _DECORATOR_SYNTAX_FLAG: [],
+}
+
+
+def _flag_enabled(flag: str) -> bool:
+    """Return true if the flag is set, or any of its superseding flags."""
+    return global_config.experimental_features[flag] or any(
+        global_config.experimental_features[f] for f in _SUPERSEDING_FLAGS[flag]
+    )
diff --git a/src/hera/workflows/_meta_mixins.py b/src/hera/workflows/_meta_mixins.py
@@ -12,6 +12,7 @@
 from typing_extensions import ParamSpec
 
 from hera.shared import BaseMixin, global_config
+from hera.shared._global_config import _DECORATOR_SYNTAX_FLAG, _flag_enabled
 from hera.shared._pydantic import BaseModel, get_fields, root_validator
 from hera.workflows._context import _context
 from hera.workflows.exceptions import InvalidTemplateCall
@@ -533,14 +534,27 @@ def create_subnode(
     return subnode
 
 
-_DECORATOR_SYNTAX_FLAG = "decorator_syntax"
-
-
 class TemplateDecoratorFuncsMixin(ContextMixin):
     from hera.workflows.dag import DAG
     from hera.workflows.script import Script
     from hera.workflows.steps import Steps
 
+    @staticmethod
+    def _check_if_enabled(decorator_name: str):
+        if not _flag_enabled(_DECORATOR_SYNTAX_FLAG):
+            raise ValueError(
+                str(
+                    "Unable to use {} decorator since it is an experimental feature."
+                    " Please turn on experimental features by setting "
+                    '`hera.shared.global_config.experimental_features["{}"] = True`.'
+                    " Note that experimental features are unstable and subject to breaking changes."
+                ).format(decorator_name, _DECORATOR_SYNTAX_FLAG)
+            )
+        if not _varname_imported:
+            raise ImportError(
+                "`varname` is not installed. Install `hera[experimental]` to bring in the extra dependency"
+            )
+
     @_add_type_hints(Script)  # type: ignore
     def script(self, **script_kwargs) -> Callable:
         """A decorator that wraps a function into a Script object.
@@ -561,19 +575,7 @@ def script(self, **script_kwargs) -> Callable:
             Function wrapper that holds a `Script` and allows the function to be called to create a Step or Task if
             in a Steps or DAG context.
         """
-        if not global_config.experimental_features[_DECORATOR_SYNTAX_FLAG]:
-            raise ValueError(
-                str(
-                    "Unable to use {} decorator since it is an experimental feature."
-                    " Please turn on experimental features by setting "
-                    '`hera.shared.global_config.experimental_features["{}"] = True`.'
-                    " Note that experimental features are unstable and subject to breaking changes."
-                ).format("script", _DECORATOR_SYNTAX_FLAG)
-            )
-        if not _varname_imported:
-            raise ImportError(
-                "`varname` is not installed. Install `hera[experimental]` to bring in the extra dependency"
-            )
+        self._check_if_enabled("script")
 
         from hera.workflows.script import RunnerScriptConstructor, Script
 
@@ -772,39 +774,15 @@ def call_wrapper(*args, **kwargs):
 
     @_add_type_hints(DAG)  # type: ignore
     def dag(self, **dag_kwargs) -> Callable:
-        if not global_config.experimental_features[_DECORATOR_SYNTAX_FLAG]:
-            raise ValueError(
-                str(
-                    "Unable to use {} decorator since it is an experimental feature."
-                    " Please turn on experimental features by setting "
-                    '`hera.shared.global_config.experimental_features["{}"] = True`.'
-                    " Note that experimental features are unstable and subject to breaking changes."
-                ).format("dag", _DECORATOR_SYNTAX_FLAG)
-            )
-        if not _varname_imported:
-            raise ImportError(
-                "`varname` is not installed. Install `hera[experimental]` to bring in the extra dependency"
-            )
+        self._check_if_enabled("dag")
 
         from hera.workflows.dag import DAG
 
         return self._construct_invocator_decorator(DAG, **dag_kwargs)
 
     @_add_type_hints(Steps)  # type: ignore
     def steps(self, **steps_kwargs) -> Callable:
-        if not global_config.experimental_features[_DECORATOR_SYNTAX_FLAG]:
-            raise ValueError(
-                str(
-                    "Unable to use {} decorator since it is an experimental feature."
-                    " Please turn on experimental features by setting "
-                    '`hera.shared.global_config.experimental_features["{}"] = True`.'
-                    " Note that experimental features are unstable and subject to breaking changes."
-                ).format("steps", _DECORATOR_SYNTAX_FLAG)
-            )
-        if not _varname_imported:
-            raise ImportError(
-                "`varname` is not installed. Install `hera[experimental]` to bring in the extra dependency"
-            )
+        self._check_if_enabled("steps")
 
         from hera.workflows.steps import Steps
 

diff --git a/src/hera/workflows/_mixins.py b/src/hera/workflows/_mixins.py
@@ -712,12 +712,6 @@ class TemplateInvocatorSubNodeMixin(BaseMixin):
 
     _build_obj: Optional[HeraBuildObj] = PrivateAttr(None)
 
-    def __init__(self, **kwargs) -> None:
-        if _context.declaring:
-            object.__setattr__(self, "_build_data", kwargs)
-        else:
-            super().__init__(**kwargs)
-
     def __getattribute__(self, name: str) -> Any:
         if _context.declaring:
             # Use object's __getattribute__ to avoid infinite recursion

diff --git a/src/hera/workflows/script.py b/src/hera/workflows/script.py
@@ -27,6 +27,11 @@
 
 from hera.expr import g
 from hera.shared import BaseMixin, global_config
+from hera.shared._global_config import (
+    _SCRIPT_ANNOTATIONS_FLAG,
+    _SCRIPT_PYDANTIC_IO_FLAG,
+    _flag_enabled,
+)
 from hera.shared._pydantic import _PYDANTIC_VERSION, root_validator, validator
 from hera.workflows._context import _context
 from hera.workflows._meta_mixins import CallableTemplateMixin
@@ -254,7 +259,7 @@ def _build_inputs(self) -> Optional[ModelInputs]:
         func_parameters: List[Parameter] = []
         func_artifacts: List[Artifact] = []
         if callable(self.source):
-            if global_config.experimental_features["script_annotations"]:
+            if _flag_enabled(_SCRIPT_ANNOTATIONS_FLAG):
                 func_parameters, func_artifacts = _get_inputs_from_callable(self.source)
             else:
                 func_parameters = _get_parameters_from_callable(self.source)
@@ -267,7 +272,7 @@ def _build_outputs(self) -> Optional[ModelOutputs]:
         if not callable(self.source):
             return outputs
 
-        if not global_config.experimental_features["script_annotations"]:
+        if not _flag_enabled(_SCRIPT_ANNOTATIONS_FLAG):
             return outputs
 
         outputs_directory = None
@@ -385,14 +390,14 @@ def append_annotation(annotation: Union[Artifact, Parameter]):
 
             append_annotation(get_args(annotation)[1])
     elif return_annotation and issubclass(return_annotation, (OutputV1, OutputV2)):
-        if not global_config.experimental_features["script_pydantic_io"]:
+        if not _flag_enabled(_SCRIPT_PYDANTIC_IO_FLAG):
             raise ValueError(
                 (
                     "Unable to instantiate {} since it is an experimental feature."
                     " Please turn on experimental features by setting "
                     '`hera.shared.global_config.experimental_features["{}"] = True`.'
                     " Note that experimental features are unstable and subject to breaking changes."
-                ).format(return_annotation, "script_pydantic_io")
+                ).format(return_annotation, _SCRIPT_PYDANTIC_IO_FLAG)
             )
 
         output_class = return_annotation
@@ -454,14 +459,14 @@ class will be used as inputs, rather than the class itself.
 
     for func_param in inspect.signature(source).parameters.values():
         if get_origin(func_param.annotation) is None and issubclass(func_param.annotation, (InputV1, InputV2)):
-            if not global_config.experimental_features["script_pydantic_io"]:
+            if not _flag_enabled(_SCRIPT_PYDANTIC_IO_FLAG):
                 raise ValueError(
                     (
                         "Unable to instantiate {} since it is an experimental feature."
                         " Please turn on experimental features by setting "
                         '`hera.shared.global_config.experimental_features["{}"] = True`.'
                         " Note that experimental features are unstable and subject to breaking changes."
-                    ).format(func_param.annotation, "script_pydantic_io")
+                    ).format(func_param.annotation, _SCRIPT_PYDANTIC_IO_FLAG)
                 )
 
             if len(inspect.signature(source).parameters) != 1:
@@ -565,7 +570,7 @@ def _output_annotations_used(source: Callable) -> bool:
 
     This includes parameters marked as outputs and the return annotation of `source`.
     """
-    if not global_config.experimental_features["script_annotations"]:
+    if not _flag_enabled(_SCRIPT_ANNOTATIONS_FLAG):
         return False
     if not callable(source):
         return False
@@ -807,15 +812,15 @@ def transform_script_template_post_build(
         self, instance: "Script", script: _ModelScriptTemplate
     ) -> _ModelScriptTemplate:
         """A hook to transform the generated script template."""
-        if global_config.experimental_features["script_annotations"]:
+        if _flag_enabled(_SCRIPT_ANNOTATIONS_FLAG):
             if not script.env:
                 script.env = []
             script.env.append(EnvVar(name="hera__script_annotations", value=""))
             if self.outputs_directory:
                 script.env.append(EnvVar(name="hera__outputs_directory", value=self.outputs_directory))
             if self.pydantic_mode:
                 script.env.append(EnvVar(name="hera__pydantic_mode", value=str(self.pydantic_mode)))
-            if global_config.experimental_features["script_pydantic_io"]:
+            if _flag_enabled(_SCRIPT_PYDANTIC_IO_FLAG):
                 script.env.append(EnvVar(name="hera__script_pydantic_io", value=""))
         return script