[FlyteSchema] Fix numpy problems

flytekit/core/type_engine.py

@@ -1495,6 +1495,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
         is_ambiguous = False
         res = None
         res_type = None
+        t = None
         for i in range(len(get_args(python_type))):
             try:
                 t = get_args(python_type)[i]
@@ -1504,8 +1505,8 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
                 if found_res:
                     is_ambiguous = True
                 found_res = True
-            except Exception:
-                logger.debug(f"Failed to convert from {python_val} to {t}", exc_info=True)
+            except Exception as e:
+                logger.debug(f"Failed to convert from {python_val} to {t} with error: {e}", exc_info=True)
                 continue
 
         if is_ambiguous:

flytekit/types/schema/types.py

@@ -9,7 +9,6 @@
 from pathlib import Path
 from typing import Type
 
-import numpy as _np
 from dataclasses_json import config
 from marshmallow import fields
 from mashumaro.mixins.json import DataClassJSONMixin
@@ -349,27 +348,39 @@ def as_readonly(self) -> FlyteSchema:
         return s
 
 
+def _get_numpy_type_mappings() -> typing.Dict[Type, SchemaType.SchemaColumn.SchemaColumnType]:
+    try:
+        import numpy as _np
+
+        return {
+            _np.int32: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
+            _np.int64: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
+            _np.uint32: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
+            _np.uint64: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
+            _np.float32: SchemaType.SchemaColumn.SchemaColumnType.FLOAT,
+            _np.float64: SchemaType.SchemaColumn.SchemaColumnType.FLOAT,
+            _np.bool_: SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN,  # type: ignore
+            _np.datetime64: SchemaType.SchemaColumn.SchemaColumnType.DATETIME,
+            _np.timedelta64: SchemaType.SchemaColumn.SchemaColumnType.DURATION,
+            _np.bytes_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
+            _np.str_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
+            _np.object_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
+        }
+    except ImportError as e:
+        logger.warning("Numpy not found, skipping numpy type mappings, error: %s", e)
+        return {}
+
+
 class FlyteSchemaTransformer(TypeTransformer[FlyteSchema]):
     _SUPPORTED_TYPES: typing.Dict[Type, SchemaType.SchemaColumn.SchemaColumnType] = {
-        _np.int32: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
-        _np.int64: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
-        _np.uint32: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
-        _np.uint64: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
-        int: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
-        _np.float32: SchemaType.SchemaColumn.SchemaColumnType.FLOAT,
-        _np.float64: SchemaType.SchemaColumn.SchemaColumnType.FLOAT,
         float: SchemaType.SchemaColumn.SchemaColumnType.FLOAT,
-        _np.bool_: SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN,  # type: ignore
+        int: SchemaType.SchemaColumn.SchemaColumnType.INTEGER,
         bool: SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN,
-        _np.datetime64: SchemaType.SchemaColumn.SchemaColumnType.DATETIME,
+        str: SchemaType.SchemaColumn.SchemaColumnType.STRING,
         datetime.datetime: SchemaType.SchemaColumn.SchemaColumnType.DATETIME,
-        _np.timedelta64: SchemaType.SchemaColumn.SchemaColumnType.DURATION,
         datetime.timedelta: SchemaType.SchemaColumn.SchemaColumnType.DURATION,
-        _np.bytes_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
-        _np.str_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
-        _np.object_: SchemaType.SchemaColumn.SchemaColumnType.STRING,
-        str: SchemaType.SchemaColumn.SchemaColumnType.STRING,
     }
+    _SUPPORTED_TYPES.update(_get_numpy_type_mappings())
 
     def __init__(self):
         super().__init__("FlyteSchema Transformer", FlyteSchema)

tests/flytekit/unit/core/test_dataclass.py

@@ -2,6 +2,7 @@
 from dataclasses_json import DataClassJsonMixin
 from mashumaro.mixins.json import DataClassJSONMixin
 import os
+import sys
 import tempfile
 from dataclasses import dataclass
 from typing import Annotated, List, Dict, Optional
@@ -882,3 +883,33 @@ class NestedFlyteTypesWithDataClassJson(DataClassJsonMixin):
     transformer = DataclassTransformer()
     lt = transformer.get_literal_type(NestedFlyteTypesWithDataClassJson)
     assert lt.metadata is not None
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="Requires Python 3.10 or higher")
+def test_numpy_import_issue_from_flyte_schema_in_dataclass():
+    from dataclasses import dataclass
+
+    from flytekit import task, workflow
+    from flytekit.types.directory import FlyteDirectory
+    from flytekit.types.file import FlyteFile
+
+    @dataclass
+    class MyDataClass:
+        output_file: FlyteFile
+        output_directory: FlyteDirectory
+
+    @task
+    def my_flyte_workflow(b: bool) -> list[MyDataClass | None]:
+        if b:
+            return [MyDataClass(__file__, ".")]
+        return [None]
+
+    @task
+    def my_flyte_task(inputs: list[MyDataClass | None]) -> bool:
+        return inputs and inputs[0] is not None
+
+    @workflow
+    def main_flyte_workflow(b: bool = False) -> bool:
+        inputs = my_flyte_workflow(b=b)
+        return my_flyte_task(inputs=inputs)
+
+    assert main_flyte_workflow(b=True) == True
+    assert main_flyte_workflow(b=False) == False