diff --git a/src/hera/workflows/script.py b/src/hera/workflows/script.py index 9a44db23b..2c817a09d 100644 --- a/src/hera/workflows/script.py +++ b/src/hera/workflows/script.py @@ -378,6 +378,18 @@ def _get_parameters_from_callable(source: Callable) -> List[Parameter]: return parameters +def _assert_pydantic_io_enabled(annotation: str) -> None: + 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(annotation, _SCRIPT_PYDANTIC_IO_FLAG) + ) + + def _get_outputs_from_return_annotation( source: Callable, outputs_directory: Optional[str], @@ -407,16 +419,7 @@ def append_annotation(annotation: Union[Artifact, Parameter]): if param_or_artifact := get_workflow_annotation(annotation): append_annotation(param_or_artifact) elif return_annotation and issubclass(return_annotation, (OutputV1, OutputV2)): - 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_FLAG) - ) - + _assert_pydantic_io_enabled(return_annotation) output_class = return_annotation for output in output_class._get_outputs(): append_annotation(output) @@ -470,15 +473,7 @@ class will be used as inputs, rather than the class itself. for func_param in inspect.signature(source).parameters.values(): if not is_subscripted(func_param.annotation) and issubclass(func_param.annotation, (InputV1, InputV2)): - 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_FLAG) - ) + _assert_pydantic_io_enabled(func_param.annotation) if len(inspect.signature(source).parameters) != 1: raise SyntaxError("Only one function parameter can be specified when using an Input.")