refactor: format and lint

-e
Signed-off-by: machichima <nary12321@gmail.com>

Author: machichima

flytekit/core/array_node_map_task.py

@@ -17,7 +17,6 @@
 from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager
 from flytekit.core.interface import transform_interface_to_list_interface
 from flytekit.core.launch_plan import LaunchPlan
-from flytekit.core.promise import Promise, create_native_named_tuple, create_task_output
 from flytekit.core.python_function_task import PythonFunctionTask, PythonInstanceTask
 from flytekit.core.task import ReferenceTask
 from flytekit.core.type_engine import TypeEngine

flytekit/core/base_task.py

@@ -610,16 +610,26 @@ def _literal_map_to_python_input(
     ) -> Dict[str, Any]:
         return TypeEngine.literal_map_to_kwargs(ctx, literal_map, self.python_interface.inputs)
 
-    async def _output_to_literal_map(self, native_outputs: Union[Tuple[Any], List[Any], Dict[str, Any]], ctx: FlyteContext):
+    async def _output_to_literal_map(
+        self, native_outputs: Union[Tuple[Any], List[Any], Dict[str, Any]], ctx: FlyteContext
+    ) -> Tuple[_literal_models.LiteralMap, Dict[str, Any]]:
         """
-        Convert output to LiteralMap.
+        Convert task output to a Flyte LiteralMap.
 
         Args:
-            native_outputs: The output from task. Can be:
-                1. Tuple[Any]: Outputs from normal tasks
-                2. List[Any]: Outputs from array node map tasks
-                3. Dict[str, Any]: Outputs from container tasks
-            ctx: The Flyte context object
+            native_outputs (Union[Tuple[Any], List[Any], Dict[str, Any]]): The output from the task. Can be:
+                1. Tuple[Any]: Outputs from normal tasks.
+                2. List[Any]: Outputs from array node map tasks.
+                3. Dict[str, Any]: Outputs from container tasks.
+            ctx (FlyteContext): The Flyte context object.
+
+        Returns:
+            Tuple[_literal_models.LiteralMap, Dict[str, Any]]: A tuple containing the converted LiteralMap and the
+            native outputs mapped by output variable names.
+
+        Raises:
+            TypeError: If an output value is a tuple where it shouldn't be.
+            BaseException: If conversion to a literal fails. The original exception is re-raised with additional context.
         """
         expected_output_names = list(self._outputs_interface.keys())
         if len(expected_output_names) == 1:
@@ -640,7 +650,9 @@ async def _output_to_literal_map(self, native_outputs: Union[Tuple[Any], List[An
             if isinstance(native_outputs, dict):
                 native_outputs_as_map = native_outputs
             else:
-                native_outputs_as_map = {expected_output_names[i]: native_outputs[i] for i, _ in enumerate(native_outputs)}
+                native_outputs_as_map = {
+                    expected_output_names[i]: native_outputs[i] for i, _ in enumerate(native_outputs)
+                }
 
         # We manually construct a LiteralMap here because task inputs and outputs actually violate the assumption
         # built into the IDL that all the values of a literal map are of the same type.
@@ -669,7 +681,8 @@ async def _output_to_literal_map(self, native_outputs: Union[Tuple[Any], List[An
                         f"Failed to convert type {type(native_outputs_as_map[expected_output_names[i]])} to type {py_type}.\n"
                         f"Error Message: {e.args[0]}.",
                     )
-                    raise e
+                    raise BaseException(*e.args) from e
+
                 literals[k2] = v2.result()
 
             if omt is not None:

flytekit/core/container_task.py

@@ -260,8 +260,6 @@ def execute(self, **kwargs) -> LiteralMap:
         except ImportError:
             raise ImportError(DOCKER_IMPORT_ERROR_MESSAGE)
 
-        from flytekit.core.type_engine import TypeEngine
-
         ctx = FlyteContext.current_context()
 
         # Normalize the input and output directories