From 90450d1d481488a858e264f082df9baa57ec7585 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Thu, 19 Dec 2024 00:38:37 +0800 Subject: [PATCH] Fix pydantic default input Signed-off-by: Future-Outlier --- flytekit/clis/sdk_in_container/run.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 7d661c3ff8..e50c4cda1d 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -53,6 +53,7 @@ labels_callback, ) from flytekit.interaction.string_literals import literal_string_repr +from flytekit.lazy_import.lazy_module import is_imported from flytekit.loggers import logger from flytekit.models import security from flytekit.models.common import RawOutputDataConfig @@ -475,8 +476,21 @@ def to_click_option( If no custom logic exists, fall back to json.dumps. """ with FlyteContextManager.with_context(flyte_ctx.new_builder()): - encoder = JSONEncoder(python_type) - default_val = encoder.encode(default_val) + if is_imported("pydantic"): + try: + from pydantic import BaseModel as BaseModelV2 + from pydantic.v1 import BaseModel as BaseModelV1 + + if issubclass(python_type, BaseModelV2): + default_val = default_val.model_dump_json() + elif issubclass(python_type, BaseModelV1): + default_val = default_val.json() + except ImportError: + # Pydantic BaseModel v1 + default_val = default_val.json() + else: + encoder = JSONEncoder(python_type) + default_val = encoder.encode(default_val) if literal_var.type.metadata: description_extra = f": {json.dumps(literal_var.type.metadata)}"