Fix Pydantic deserialization for FlyteFile and FlyteDirectory

  Fixes #6669

Signed-off-by: Govert Verkes <[email protected]>

Author: gverkes

flytekit/types/directory/types.py

@@ -179,7 +179,17 @@ def serialize_flyte_dir(self) -> Dict[str, str]:
     @model_validator(mode="after")
     def deserialize_flyte_dir(self, info) -> FlyteDirectory:
         if info.context is None or info.context.get("deserialize") is not True:
-            return self
+            # Check if all private attributes are already set up (e.g., from __init__)
+            if hasattr(self, "_downloader") and hasattr(self, "_remote_source"):
+                return self
+
+            # Populate missing private attributes for Pydantic-deserialized instances
+            dict_obj = {"path": str(self.path)}
+
+            return FlyteDirToMultipartBlobTransformer().dict_to_flyte_directory(
+                dict_obj=dict_obj,
+                expected_python_type=type(self),
+            )
 
         pv = FlyteDirToMultipartBlobTransformer().to_python_value(
             FlyteContextManager.current_context(),

flytekit/types/file/file.py

@@ -196,7 +196,18 @@ def serialize_flyte_file(self) -> Dict[str, typing.Any]:
     @model_validator(mode="after")
     def deserialize_flyte_file(self, info) -> "FlyteFile":
         if info.context is None or info.context.get("deserialize") is not True:
-            return self
+            if hasattr(self, "_downloader") and hasattr(self, "_remote_source"):
+                return self
+
+            dict_obj = {"path": str(self.path)}
+            metadata = getattr(self, "metadata", None)
+            if metadata is not None:
+                dict_obj["metadata"] = metadata
+
+            return FlyteFilePathTransformer().dict_to_flyte_file(
+                dict_obj=dict_obj,
+                expected_python_type=type(self),
+            )
 
         pv = FlyteFilePathTransformer().to_python_value(
             FlyteContextManager.current_context(),

tests/flytekit/unit/extras/pydantic_transformer/test_pydantic_basemodel_transformer.py

@@ -1022,3 +1022,37 @@ def mock_resolve_remote_path(flyte_uri: str):
 
     bm_revived = TypeEngine.to_python_value(ctx, lit, BM)
     assert bm_revived.s.literal.uri == "/my/replaced/val"
+
+
+def test_flytefile_pydantic_model_dump_validate_cycle():
+    class BM(BaseModel):
+        ff: FlyteFile
+
+    bm = BM(ff=FlyteFile.from_source("s3://my-bucket/file.txt"))
+
+    assert bm.ff.remote_source == "s3://my-bucket/file.txt"
+
+    bm_dict = bm.model_dump()
+    bm2 = BM.model_validate(bm_dict)
+
+    assert isinstance(bm2.ff, FlyteFile)
+    assert bm2.ff.remote_source == "s3://my-bucket/file.txt"
+
+    bm2.model_dump()
+
+
+def test_flytedirectory_pydantic_model_dump_validate_cycle():
+    class BM(BaseModel):
+        fd: FlyteDirectory
+
+    bm = BM(fd=FlyteDirectory.from_source("s3://my-bucket/my-dir"))
+
+    assert bm.fd.remote_source == "s3://my-bucket/my-dir"
+
+    bm_dict = bm.model_dump()
+    bm2 = BM.model_validate(bm_dict)
+
+    assert isinstance(bm2.fd, FlyteDirectory)
+    assert bm2.fd.remote_source == "s3://my-bucket/my-dir"
+
+    bm2.model_dump()