clean up if statements

Signed-off-by: Daniel Sola <[email protected]>
flyteorg · Oct 10, 2024 · 8e42042 · 8e42042
1 parent 1c3733e
commit 8e42042
Showing 1 changed file with 16 additions and 26 deletions.
diff --git a/flytekit/tools/repo.py b/flytekit/tools/repo.py
@@ -209,7 +209,8 @@ def secho(i: Identifier, state: str = "success", reason: str = None, op: str = "
         state_ind = "\r[✔]"
         fg = "green"
         nl = True
-        reason = f"successful: {console_url}" if not reason else reason
+        if not reason:
+            reason = f"successful: {console_url}" if console_url else f"version: {i.version}"
     elif state == "failed":
         state_ind = "\r[x]"
         fg = "red"
@@ -305,6 +306,11 @@ def register(
         return
 
     def _raw_register(cp_entity: FlyteControlPlaneEntity):
+        fetch_methods = {
+            task.TaskSpec: remote.fetch_task,
+            admin_workflow_models.WorkflowSpec: remote.fetch_workflow,
+            launch_plan_models.LaunchPlan: remote.fetch_launch_plan,
+        }
         is_lp = False
         if isinstance(cp_entity, launch_plan.LaunchPlan):
             og_id = cp_entity.id
@@ -317,32 +323,16 @@ def _raw_register(cp_entity: FlyteControlPlaneEntity):
                     i = remote.raw_register(
                         cp_entity, serialization_settings, version=version, create_default_launchplan=False
                     )
-                    # Get URL
+                    # Get the fetch method based on the type of cp_entity
+                    fetch_method = next(
+                        (method for entity_type, method in fetch_methods.items() if isinstance(cp_entity, entity_type)),
+                        None,
+                    )
+                    # If a valid fetch method was found, fetch the entity and generate the URL
                     console_url = None
-                    if isinstance(cp_entity, task.TaskSpec):
-                        task_entity = remote.fetch_task(
-                            i.project,
-                            i.domain,
-                            i.name,
-                            i.version,
-                        )
-                        console_url = remote.generate_console_url(task_entity)
-                    if isinstance(cp_entity, admin_workflow_models.WorkflowSpec):
-                        workflow_entity = remote.fetch_workflow(
-                            i.project,
-                            i.domain,
-                            i.name,
-                            i.version,
-                        )
-                        console_url = remote.generate_console_url(workflow_entity)
-                    if isinstance(cp_entity, launch_plan_models.LaunchPlan):
-                        workflow_entity = remote.fetch_launch_plan(
-                            i.project,
-                            i.domain,
-                            i.name,
-                            i.version,
-                        )
-                        console_url = remote.generate_console_url(workflow_entity)
+                    if fetch_method:
+                        entity = fetch_method(i.project, i.domain, i.name, i.version)
+                        console_url = remote.generate_console_url(entity)
 
                     secho(i, state="success", console_url=console_url)
                     if is_lp and activate_launchplans: