From f2f1a7de513b9be2127c2cc3ff922d49f289a380 Mon Sep 17 00:00:00 2001
From: connor-mccarthy <mccarthy.connor.james@gmail.com>
Date: Wed, 2 Nov 2022 11:36:43 -0700
Subject: [PATCH] clean up typing

---
 sdk/python/kfp/components/pipeline_task.py | 4 ++--
 sdk/python/kfp/components/placeholders.py  | 6 +-----
 sdk/python/kfp/components/structures.py    | 2 +-
 3 files changed, 4 insertions(+), 8 deletions(-)

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):