From eb27e374260bbde00b51c6e265d92ac6c751261d Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Fri, 15 Nov 2024 10:26:53 +0800 Subject: [PATCH] Union Dataclass Ambiguos Error Signed-off-by: Future-Outlier --- flytekit/core/type_engine.py | 50 ++++++++------- .../core/test_generice_idl_type_engine.py | 63 +++++++++++-------- tests/flytekit/unit/core/test_type_engine.py | 63 +++++++++++-------- 3 files changed, 98 insertions(+), 78 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 160c0595a4..846c8e372c 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -608,36 +608,38 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: schema = None try: - from marshmallow_enum import EnumField, LoadDumpOptions - - if issubclass(t, DataClassJsonMixin): - s = cast(DataClassJsonMixin, self._get_origin_type_in_annotation(t)).schema() - for _, v in s.fields.items(): - # marshmallow-jsonschema only supports enums loaded by name. - # https://github.com/fuhrysteve/marshmallow-jsonschema/blob/81eada1a0c42ff67de216923968af0a6b54e5dcb/marshmallow_jsonschema/base.py#L228 - if isinstance(v, EnumField): - v.load_by = LoadDumpOptions.name - # check if DataClass mixin - from marshmallow_jsonschema import JSONSchema - - schema = JSONSchema().dump(s) + # This produce JSON SCHEMA draft 2020-12 + from mashumaro.jsonschema import build_json_schema + + schema = build_json_schema(cast(DataClassJSONMixin, self._get_origin_type_in_annotation(t))).to_dict() except Exception as e: - # https://github.com/lovasoa/marshmallow_dataclass/issues/13 - logger.warning( - f"Failed to extract schema for object {t}, (will run schemaless) error: {e}" - f"If you have postponed annotations turned on (PEP 563) turn it off please. Postponed" - f"evaluation doesn't work with json dataclasses" + logger.error( + f"Failed to extract schema for object {t}, error: {e}\n" + f"Please remove `DataClassJsonMixin` and `dataclass_json` decorator from the dataclass definition" ) if schema is None: try: - from mashumaro.jsonschema import build_json_schema - - schema = build_json_schema(cast(DataClassJSONMixin, self._get_origin_type_in_annotation(t))).to_dict() + # This produce JSON SCHEMA draft 2020-12 + from marshmallow_enum import EnumField, LoadDumpOptions + + if issubclass(t, DataClassJsonMixin): + s = cast(DataClassJsonMixin, self._get_origin_type_in_annotation(t)).schema() + for _, v in s.fields.items(): + # marshmallow-jsonschema only supports enums loaded by name. + # https://github.com/fuhrysteve/marshmallow-jsonschema/blob/81eada1a0c42ff67de216923968af0a6b54e5dcb/marshmallow_jsonschema/base.py#L228 + if isinstance(v, EnumField): + v.load_by = LoadDumpOptions.name + # check if DataClass mixin + from marshmallow_jsonschema import JSONSchema + + schema = JSONSchema().dump(s) except Exception as e: - logger.error( - f"Failed to extract schema for object {t}, error: {e}\n" - f"Please remove `DataClassJsonMixin` and `dataclass_json` decorator from the dataclass definition" + # https://github.com/lovasoa/marshmallow_dataclass/issues/13 + logger.warning( + f"Failed to extract schema for object {t}, (will run schemaless) error: {e}" + f"If you have postponed annotations turned on (PEP 563) turn it off please. Postponed" + f"evaluation doesn't work with json dataclasses" ) # Recursively construct the dataclass_type which contains the literal type of each field diff --git a/tests/flytekit/unit/core/test_generice_idl_type_engine.py b/tests/flytekit/unit/core/test_generice_idl_type_engine.py index f23c9557f8..3f45ee9f27 100644 --- a/tests/flytekit/unit/core/test_generice_idl_type_engine.py +++ b/tests/flytekit/unit/core/test_generice_idl_type_engine.py @@ -751,39 +751,48 @@ class UnsupportedNestedStruct(DataClassJsonMixin): s: UnsupportedSchemaType schema = { - "$ref": "#/definitions/TeststructSchema", - "$schema": "http://json-schema.org/draft-07/schema#", - "definitions": { - "InnerstructSchema": { - "additionalProperties": False, - "properties": { - "a": {"title": "a", "type": "integer"}, - "b": {"default": None, "title": "b", "type": ["string", "null"]}, - "c": { - "items": {"title": "c", "type": "integer"}, - "title": "c", - "type": "array", - }, - }, + "type": "object", + "title": "TestStruct", + "properties": { + "s": { "type": "object", - }, - "TeststructSchema": { - "additionalProperties": False, + "title": "InnerStruct", "properties": { - "m": { - "additionalProperties": {"title": "m", "type": "string"}, - "title": "m", - "type": "object", + "a": { + "type": "integer" }, - "s": { - "$ref": "#/definitions/InnerstructSchema", - "field_many": False, - "type": "object", + "b": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] }, + "c": { + "type": "array", + "items": { + "type": "integer" + } + } }, - "type": "object", + "additionalProperties": False, + "required": ["a", "b", "c"] }, + "m": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "type": "string" + } + } }, + "additionalProperties": False, + "required": ["s", "m"] } tf = DataclassTransformer() t = tf.get_literal_type(TestStruct) @@ -1566,7 +1575,7 @@ class Bar(DataClassJsonMixin): pv = Bar(x=3) with pytest.raises( TypeTransformerFailedError, - match="Type of Val '' is not an instance of ", + match="Type of Val '' is not an instance of ", ): DataclassTransformer().assert_type(gt, pv) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 1d6552a7c5..68552a1848 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -741,39 +741,48 @@ class UnsupportedNestedStruct(DataClassJsonMixin): s: UnsupportedSchemaType schema = { - "$ref": "#/definitions/TeststructSchema", - "$schema": "http://json-schema.org/draft-07/schema#", - "definitions": { - "InnerstructSchema": { - "additionalProperties": False, - "properties": { - "a": {"title": "a", "type": "integer"}, - "b": {"default": None, "title": "b", "type": ["string", "null"]}, - "c": { - "items": {"title": "c", "type": "integer"}, - "title": "c", - "type": "array", - }, - }, + "type": "object", + "title": "TestStruct", + "properties": { + "s": { "type": "object", - }, - "TeststructSchema": { - "additionalProperties": False, + "title": "InnerStruct", "properties": { - "m": { - "additionalProperties": {"title": "m", "type": "string"}, - "title": "m", - "type": "object", + "a": { + "type": "integer" }, - "s": { - "$ref": "#/definitions/InnerstructSchema", - "field_many": False, - "type": "object", + "b": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] }, + "c": { + "type": "array", + "items": { + "type": "integer" + } + } }, - "type": "object", + "additionalProperties": False, + "required": ["a", "b", "c"] }, + "m": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "type": "string" + } + } }, + "additionalProperties": False, + "required": ["s", "m"] } tf = DataclassTransformer() t = tf.get_literal_type(TestStruct) @@ -1569,7 +1578,7 @@ class Bar(DataClassJsonMixin): pv = Bar(x=3) with pytest.raises( TypeTransformerFailedError, - match="Type of Val '' is not an instance of ", + match="Type of Val '' is not an instance of ", ): DataclassTransformer().assert_type(gt, pv)