Add support FlyteFile in dataclass

flytekit/core/type_engine.py

@@ -28,7 +28,16 @@
 from flytekit.models import interface as _interface_models
 from flytekit.models import types as _type_models
 from flytekit.models.core import types as _core_types
-from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Schema
+from flytekit.models.literals import (
+    Blob,
+    BlobMetadata,
+    Literal,
+    LiteralCollection,
+    LiteralMap,
+    Primitive,
+    Scalar,
+    Schema,
+)
 from flytekit.models.types import LiteralType, SimpleType
 
 T = typing.TypeVar("T")
@@ -258,29 +267,79 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]):
         """
         If any field inside the dataclass is flyte type, we should use flyte type transformer for that field.
         """
-        from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer
+        from flytekit.types.directory.types import FlyteDirectory
+        from flytekit.types.file import FlyteFile
+        from flytekit.types.schema.types import FlyteSchema
 
         for f in dataclasses.fields(python_type):
             v = python_val.__getattribute__(f.name)
-            if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema):
-                FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None)
+            if inspect.isclass(f.type) and (
+                issubclass(f.type, FlyteSchema) or issubclass(f.type, FlyteFile) or issubclass(f.type, FlyteDirectory)
+            ):
+                TypeEngine.to_literal(FlyteContext.current_context(), v, f.type, None)
             elif dataclasses.is_dataclass(f.type):
                 self._serialize_flyte_type(v, f.type)
 
-    def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]):
+    def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> T:
+        from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer
+        from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer
         from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer
 
-        for f in dataclasses.fields(expected_python_type):
-            v = python_val.__getattribute__(f.name)
-            if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema):
-                t = FlyteSchemaTransformer()
-                t.to_python_value(
-                    FlyteContext.current_context(),
-                    Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))),
-                    f.type,
-                )
-            elif dataclasses.is_dataclass(f.type):
-                self._deserialize_flyte_type(v, f.type)
+        if not dataclasses.is_dataclass(expected_python_type):
+            return python_val
+
+        if issubclass(expected_python_type, FlyteSchema):
+            t = FlyteSchemaTransformer()
+            return t.to_python_value(
+                FlyteContext.current_context(),
+                Literal(scalar=Scalar(schema=Schema(python_val.remote_path, t._get_schema_type(expected_python_type)))),
+                expected_python_type,
+            )
+        elif issubclass(expected_python_type, FlyteFile):
+            return FlyteFilePathTransformer().to_python_value(
+                FlyteContext.current_context(),
+                Literal(
+                    scalar=Scalar(
+                        blob=Blob(
+                            metadata=BlobMetadata(
+                                type=_core_types.BlobType(
+                                    format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
+                                )
+                            ),
+                            uri=python_val.path,
+                        )
+                    )
+                ),
+                expected_python_type,
+            )
+        elif issubclass(expected_python_type, FlyteDirectory):
+            return FlyteDirToMultipartBlobTransformer().to_python_value(
+                FlyteContext.current_context(),
+                Literal(
+                    scalar=Scalar(
+                        blob=Blob(
+                            metadata=BlobMetadata(
+                                type=_core_types.BlobType(
+                                    format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART
+                                )
+                            ),
+                            uri=python_val.path,
+                        )
+                    )
+                ),
+                expected_python_type,
+            )
+        else:
+            for f in dataclasses.fields(expected_python_type):
+                value = python_val.__getattribute__(f.name)
+                if hasattr(f.type, "__origin__") and f.type.__origin__ is list:
+                    value = [self._deserialize_flyte_type(v, f.type.__args__[0]) for v in value]
+                elif hasattr(f.type, "__origin__") and f.type.__origin__ is dict:
+                    value = {k: self._deserialize_flyte_type(v, f.type.__args__[1]) for k, v in value.items()}
+                else:
+                    value = self._deserialize_flyte_type(value, f.type)
+                python_val.__setattr__(f.name, value)
+            return python_val
 
     def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any:
         if t == int:
@@ -326,8 +385,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
             )
 
         dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic))
-        self._deserialize_flyte_type(dc, expected_python_type)
-        return self._fix_dataclass_int(expected_python_type, dc)
+        return self._fix_dataclass_int(expected_python_type, self._deserialize_flyte_type(dc, expected_python_type))
 
     def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
         if literal_type.simple == SimpleType.STRUCT:
@@ -817,18 +875,21 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac
     """
     attribute_list = []
     for property_key, property_val in schema[schema_name]["properties"].items():
+        property_type = property_val["type"]
         # Handle list
         if property_val["type"] == "array":
             attribute_list.append((property_key, typing.List[_get_element_type(property_val["items"])]))
         # Handle dataclass and dict
-        elif property_val["type"] == "object":
-            if "$ref" in property_val:
+        elif property_type == "object":
+            if property_val.get("$ref"):
                 name = property_val["$ref"].split("/")[-1]
                 attribute_list.append((property_key, convert_json_schema_to_python_class(schema, name)))
-            else:
+            elif property_val.get("additionalProperties"):
                 attribute_list.append(
                     (property_key, typing.Dict[str, _get_element_type(property_val["additionalProperties"])])
                 )
+            else:
+                attribute_list.append((property_key, typing.Dict[str, _get_element_type(property_val)]))
         # Handle int, float, bool or str
         else:
             attribute_list.append([property_key, _get_element_type(property_val)])

flytekit/types/directory/types.py

@@ -3,8 +3,12 @@
 import os
 import pathlib
 import typing
+from dataclasses import dataclass, field
 from pathlib import Path
 
+from dataclasses_json import config, dataclass_json
+from marshmallow import fields
+
 from flytekit.core.context_manager import FlyteContext
 from flytekit.core.type_engine import TypeEngine, TypeTransformer
 from flytekit.models import types as _type_models
@@ -19,7 +23,10 @@ def noop():
     ...
 
 
+@dataclass_json
+@dataclass
 class FlyteDirectory(os.PathLike, typing.Generic[T]):
+    path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String()))
     """
     .. warning::
 
@@ -106,15 +113,16 @@ def t1(in1: FlyteDirectory["svg"]):
     field in the ``BlobType``.
     """
 
-    def __init__(self, path: str, downloader: typing.Callable = None, remote_directory=None):
+    def __init__(self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = None, remote_directory=None):
         """
         :param path: The source path that users are expected to call open() on
         :param downloader: Optional function that can be passed that used to delay downloading of the actual fil
             until a user actually calls open().
         :param remote_directory: If the user wants to return something and also specify where it should be uploaded to.
         """
-
-        self._path = path
+        # Make this field public, so that the dataclass transformer can set a value for it
+        # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298
+        self.path = path
         self._downloader = downloader or noop
         self._downloaded = False
         self._remote_directory = remote_directory
@@ -127,7 +135,7 @@ def __fspath__(self):
         if not self._downloaded:
             self._downloader()
             self._downloaded = True
-        return self._path
+        return self.path
 
     @classmethod
     def extension(cls) -> str:
@@ -159,10 +167,6 @@ def downloaded(self) -> bool:
     def remote_directory(self) -> typing.Optional[str]:
         return self._remote_directory
 
-    @property
-    def path(self) -> str:
-        return self._path
-
     @property
     def remote_source(self) -> str:
         """
@@ -175,10 +179,10 @@ def download(self) -> str:
         return self.__fspath__()
 
     def __repr__(self):
-        return self._path
+        return self.path
 
     def __str__(self):
-        return self._path
+        return self.path
 
 
 class FlyteDirToMultipartBlobTransformer(TypeTransformer[FlyteDirectory]):
@@ -238,7 +242,7 @@ def to_literal(
                 return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=python_val._remote_source)))
 
             source_path = python_val.path
-            # If the user specified the remote_path to be False, that means no matter what, do not upload. Also if the
+            # If the user specified the remote_directory to be False, that means no matter what, do not upload. Also if the
             # path given is already a remote path, say https://www.google.com, the concept of uploading to the Flyte
             # blob store doesn't make sense.
             if python_val.remote_directory is False or ctx.file_access.is_remote(source_path):

flytekit/types/file/file.py

@@ -3,6 +3,10 @@
 import os
 import pathlib
 import typing
+from dataclasses import dataclass, field
+
+from dataclasses_json import config, dataclass_json
+from marshmallow import fields
 
 from flytekit.core.context_manager import FlyteContext
 from flytekit.core.type_engine import TypeEngine, TypeTransformer
@@ -19,7 +23,10 @@ def noop():
 T = typing.TypeVar("T")
 
 
+@dataclass_json
+@dataclass
 class FlyteFile(os.PathLike, typing.Generic[T]):
+    path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String()))
     """
     Since there is no native Python implementation of files and directories for the Flyte Blob type, (like how int
     exists for Flyte's Integer type) we need to create one so that users can express that their tasks take
@@ -155,14 +162,18 @@ def extension(cls) -> str:
 
         return _SpecificFormatClass
 
-    def __init__(self, path: str, downloader: typing.Callable = noop, remote_path=None):
+    def __init__(
+        self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = noop, remote_path: os.PathLike = None
+    ):
         """
         :param path: The source path that users are expected to call open() on
         :param downloader: Optional function that can be passed that used to delay downloading of the actual fil
             until a user actually calls open().
         :param remote_path: If the user wants to return something and also specify where it should be uploaded to.
         """
-        self._path = path
+        # Make this field public, so that the dataclass transformer can set a value for it
+        # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298
+        self.path = path
         self._downloader = downloader
         self._downloaded = False
         self._remote_path = remote_path
@@ -173,30 +184,26 @@ def __fspath__(self):
         if not self._downloaded:
             self._downloader()
             self._downloaded = True
-        return self._path
+        return self.path
 
     def __eq__(self, other):
         if isinstance(other, FlyteFile):
             return (
-                self._path == other._path
+                self.path == other.path
                 and self._remote_path == other._remote_path
                 and self.extension() == other.extension()
             )
         else:
-            return self._path == other
+            return self.path == other
 
     @property
     def downloaded(self) -> bool:
         return self._downloaded
 
     @property
-    def remote_path(self) -> typing.Optional[str]:
+    def remote_path(self) -> os.PathLike:
         return self._remote_path
 
-    @property
-    def path(self) -> str:
-        return self._path
-
     @property
     def remote_source(self) -> str:
         """
@@ -209,10 +216,10 @@ def download(self) -> str:
         return self.__fspath__()
 
     def __repr__(self):
-        return self._path
+        return self.path
 
     def __str__(self):
-        return self._path
+        return self.path
 
 
 class FlyteFilePathTransformer(TypeTransformer[FlyteFile]):
@@ -316,11 +323,10 @@ def to_literal(
             return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=source_path)))
 
     def to_python_value(
-        self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile]]
+        self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike]
     ) -> FlyteFile:
 
         uri = lv.scalar.blob.uri
-
         # In this condition, we still return a FlyteFile instance, but it's a simple one that has no downloading tricks
         # Using is instead of issubclass because FlyteFile does actually subclass it
         if expected_python_type is os.PathLike: