clean up typing

kubeflow · Nov 2, 2022 · f2f1a7d · f2f1a7d
1 parent be2728c
commit f2f1a7d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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):