[Docs] MessagePack IDL, Pydantic Support, and Attribute Access

dev-requirements.in

@@ -4,6 +4,7 @@ coverage
 pre-commit
 codespell
 mock
+pydantic>2
 pytest
 mypy
 mashumaro

examples/data_types_and_io/Dockerfile

@@ -19,7 +19,7 @@ RUN python3 -m venv ${VENV}
 ENV PATH="${VENV}/bin:$PATH"
 
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
-    pip install flytekit pandas pyarrow
+    pip install flytekit pandas pyarrow pydantic>2
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
     pip install torch --index-url https://download.pytorch.org/whl/cpu
 

examples/data_types_and_io/data_types_and_io/attribute_access.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 
-from dataclasses_json import dataclass_json
 from flytekit import task, workflow
 
 
@@ -36,7 +35,6 @@ def dict_wf():
 
 
 # Directly access an attribute of a dataclass
-@dataclass_json
 @dataclass
 class Fruit:
     name: str

examples/data_types_and_io/data_types_and_io/pydantic_basemodel_attribute_access.py

@@ -0,0 +1,93 @@
+from flytekit import ImageSpec, task, workflow
+from pydantic import BaseModel
+
+image_spec = ImageSpec(
+    registry="ghcr.io/flyteorg",
+    packages=["pydantic>2"],
+)
+
+
+@task(container_image=image_spec)
+def print_message(message: str):
+    print(message)
+    return
+
+
+# Access an output list using index notation
+@task(container_image=image_spec)
+def list_task() -> list[str]:
+    return ["apple", "banana"]
+
+
+@workflow
+def list_wf():
+    items = list_task()
+    first_item = items[0]
+    print_message(message=first_item)
+
+
+# Access the output dictionary by specifying the key
+@task(container_image=image_spec)
+def dict_task() -> dict[str, str]:
+    return {"fruit": "banana"}
+
+
+@workflow
+def dict_wf():
+    fruit_dict = dict_task()
+    print_message(message=fruit_dict["fruit"])
+
+
+# Directly access an attribute of a Pydantic BaseModel
+class Fruit(BaseModel):
+    name: str
+
+
+@task(container_image=image_spec)
+def basemodel_task() -> Fruit:
+    return Fruit(name="banana")
+
+
+@workflow
+def basemodel_wf():
+    fruit_instance = basemodel_task()
+    print_message(message=fruit_instance.name)
+
+
+# Combinations of list, dict, and BaseModel also work effectively
+@task(container_image=image_spec)
+def advance_task() -> (dict[str, list[str]], list[dict[str, str]], dict[str, Fruit]):
+    return (
+        {"fruits": ["banana"]},
+        [{"fruit": "banana"}],
+        {"fruit": Fruit(name="banana")},
+    )
+
+
+@task(container_image=image_spec)
+def print_list(fruits: list[str]):
+    print(fruits)
+
+
+@task(container_image=image_spec)
+def print_dict(fruit_dict: dict[str, str]):
+    print(fruit_dict)
+
+
+@workflow
+def advanced_workflow():
+    dictionary_list, list_dict, dict_basemodel = advance_task()
+    print_message(message=dictionary_list["fruits"][0])
+    print_message(message=list_dict[0]["fruit"])
+    print_message(message=dict_basemodel["fruit"].name)
+
+    print_list(fruits=dictionary_list["fruits"])
+    print_dict(fruit_dict=list_dict[0])
+
+
+# Run the workflows locally
+if __name__ == "__main__":
+    list_wf()
+    dict_wf()
+    basemodel_wf()
+    advanced_workflow()

examples/data_types_and_io/requirements.in

@@ -3,3 +3,4 @@ torch
 tabulate
 tensorflow
 pyarrow
+pydantic>2