diff --git a/sdk/python/kfp/components/pipeline_task.py b/sdk/python/kfp/components/pipeline_task.py index 4e7572cc13d..c0e1a1ba5fe 100644 --- a/sdk/python/kfp/components/pipeline_task.py +++ b/sdk/python/kfp/components/pipeline_task.py @@ -474,7 +474,7 @@ def my_pipeline(): def check_primitive_placeholder_is_used_for_correct_io_type( inputs_dict: Dict[str, structures.InputSpec], outputs_dict: Dict[str, structures.OutputSpec], - arg: Union[placeholders.Placeholder, Any], + arg: Union[placeholders.CommandLineElement, Any], ): """Validates input/output placeholders refer to an input/output with an appropriate type for the placeholder. This should only apply to components @@ -514,7 +514,7 @@ def check_primitive_placeholder_is_used_for_correct_io_type( f'"{outputs_dict[output_name].type}" cannot be paired with ' f'{arg.__class__.__name__}.') elif isinstance(arg, placeholders.IfPresentPlaceholder): - all_normalized_args: List[Union[str, placeholders.Placeholder]] = [] + all_normalized_args: List[placeholders.CommandLineElement] = [] if arg.then is None: pass elif isinstance(arg.then, list): diff --git a/sdk/python/kfp/components/placeholders.py b/sdk/python/kfp/components/placeholders.py index 3d249073cfd..2a0f5346c3b 100644 --- a/sdk/python/kfp/components/placeholders.py +++ b/sdk/python/kfp/components/placeholders.py @@ -265,11 +265,7 @@ def _to_string(self) -> str: OutputPathPlaceholder, OutputUriPlaceholder, OutputMetadataPlaceholder) -CommandLineElement = Union[str, IfPresentPlaceholder, ConcatPlaceholder, - InputValuePlaceholder, InputPathPlaceholder, - InputUriPlaceholder, InputMetadataPlaceholder, - OutputParameterPlaceholder, OutputPathPlaceholder, - OutputUriPlaceholder, OutputMetadataPlaceholder] +CommandLineElement = Union[str, Placeholder] def convert_command_line_element_to_string( diff --git a/sdk/python/kfp/components/structures.py b/sdk/python/kfp/components/structures.py index e29a79468f4..6ec76df69bb 100644 --- a/sdk/python/kfp/components/structures.py +++ b/sdk/python/kfp/components/structures.py @@ -494,7 +494,7 @@ def check_placeholder_references_valid_io_name( f'Argument "{arg}" references nonexistant input: "{arg.input_name}".' ) - all_normalized_args: List[Union[str, placeholders.Placeholder]] = [] + all_normalized_args: List[placeholders.CommandLineElement] = [] if arg.then is None: pass elif isinstance(arg.then, list):