Merge pull request #20 from beda-software/pydantic-v2

Prepare for pydanticv2

Author: ruscoder

fhir_py_types/ast.py

@@ -1,12 +1,11 @@
 import ast
-import functools
 import itertools
 import keyword
 import logging
 from collections.abc import Iterable
 from dataclasses import replace
 from enum import Enum, auto
-from typing import Literal, cast
+from typing import Literal
 
 from fhir_py_types import (
     StructureDefinition,
@@ -79,11 +78,13 @@ def format_identifier(
     def uppercamelcase(s: str) -> str:
         return s[:1].upper() + s[1:]
 
-    return (
-        identifier + uppercamelcase(type_.code)
-        if is_polymorphic(definition)
-        else identifier
-    )
+    if is_polymorphic(definition):
+        # TODO: it's fast hack
+        if type_.code[0].islower():
+            return identifier + uppercamelcase(clear_primitive_id(type_.code))
+        return identifier + uppercamelcase(type_.code)
+
+    return identifier
 
 
 def remap_type(
@@ -94,8 +95,8 @@ def remap_type(
             case "Resource":
                 # Different contexts use 'Resource' type to refer to any
                 # resource differentiated by its 'resourceType' (tagged union).
-                # 'AnyResource' is not defined by the spec but rather
-                # generated as a union of all defined resource types.
+                # 'AnyResource' is defined in header as a special type
+                # that dynamically replaced with a right type in run-time
                 type_ = replace(type_, code="AnyResource")
 
     if is_polymorphic(definition):
@@ -106,6 +107,12 @@ def remap_type(
         # with a custom validator that will enforce single required property rule.
         type_ = replace(type_, required=False)
 
+    if is_primitive_type(type_):
+        # Primitive types defined from the small letter (like code)
+        # and it might overlap with model fields
+        # e.g. QuestionnaireItem has attribute code and linkId has type code
+        type_ = replace(type_, code=make_primitive_id(type_.code))
+
     return type_
 
 
@@ -115,8 +122,7 @@ def zip_identifier_type(
     result = []
 
     for t in [remap_type(definition, t) for t in definition.type]:
-        name = format_identifier(definition, identifier, t)
-        result.append((name, t))
+        result.append((format_identifier(definition, identifier, t), t))
         if definition.kind != StructureDefinitionKind.PRIMITIVE and is_primitive_type(
             t
         ):
@@ -185,10 +191,16 @@ def order_type_overriding_properties(
 def define_class_object(
     definition: StructureDefinition,
 ) -> Iterable[ast.stmt | ast.expr]:
+    bases: list[ast.expr] = []
+    if definition.kind == StructureDefinitionKind.RESOURCE:
+        bases.append(ast.Name("AnyResource"))
+    # BaseModel should be the last, because it overrides `extra`
+    bases.append(ast.Name("BaseModel"))
+
     return [
         ast.ClassDef(
             definition.id,
-            bases=[ast.Name("BaseModel")],
+            bases=bases,
             body=[
                 ast.Expr(value=ast.Constant(definition.docstring)),
                 *itertools.chain.from_iterable(
@@ -202,11 +214,6 @@ def define_class_object(
             keywords=[],
             type_params=[],
         ),
-        ast.Call(
-            ast.Attribute(value=ast.Name(definition.id), attr="update_forward_refs"),
-            args=[],
-            keywords=[],
-        ),
     ]
 
 
@@ -215,48 +222,22 @@ def define_class(definition: StructureDefinition) -> Iterable[ast.stmt | ast.exp
 
 
 def define_alias(definition: StructureDefinition) -> Iterable[ast.stmt]:
-    return type_annotate(definition, definition.id, AnnotationForm.TypeAlias)
-
-
-def define_tagged_union(
-    name: str, components: Iterable[StructureDefinition], distinct_by: str
-) -> ast.stmt:
-    annotation = functools.reduce(
-        lambda acc, n: ast.BinOp(left=acc, right=n, op=ast.BitOr()),
-        (cast(ast.expr, ast.Name(d.id)) for d in components),
+    # Primitive types are renamed to another name to avoid overlapping with model fields
+    return type_annotate(
+        definition, make_primitive_id(definition.id), AnnotationForm.TypeAlias
     )
 
-    return ast.Assign(
-        targets=[ast.Name(name)],
-        value=ast.Subscript(
-            value=ast.Name("Annotated_"),
-            slice=ast.Tuple(
-                elts=[
-                    annotation,
-                    ast.Call(
-                        ast.Name("Field"),
-                        args=[ast.Constant(...)],
-                        keywords=[
-                            ast.keyword(
-                                arg="discriminator", value=ast.Constant(distinct_by)
-                            ),
-                        ],
-                    ),
-                ]
-            ),
-        ),
-    )
 
+def make_primitive_id(name: str) -> str:
+    if name in ("str", "int", "float", "bool"):
+        return name
+    return f"{name}Type"
 
-def select_tagged_resources(
-    definitions: Iterable[StructureDefinition], key: str
-) -> Iterable[StructureDefinition]:
-    return (
-        definition
-        for definition in definitions
-        if definition.kind == StructureDefinitionKind.RESOURCE
-        and key in definition.elements
-    )
+
+def clear_primitive_id(name: str) -> str:
+    if name.endswith("Type"):
+        return name[:-4]
+    return name
 
 
 def select_nested_definitions(
@@ -302,14 +283,6 @@ def build_ast(
                         f"Unsupported definition {definition.id} of kind {definition.kind}, skipping"
                     )
 
-    resources = list(select_tagged_resources(structure_definitions, key="resourceType"))
-    if resources:
-        typedefinitions.append(
-            define_tagged_union(
-                name="AnyResource", components=resources, distinct_by="resourceType"
-            )
-        )
-
     return sorted(
         typedefinitions,
         # Defer any postprocessing until after the structure tree is defined.

fhir_py_types/header.py.tpl

@@ -1,15 +1,162 @@
-from typing import List as List_, Optional as Optional_, Literal as Literal_, Annotated as Annotated_
+from typing import (
+    List as List_,
+    Optional as Optional_,
+    Literal as Literal_,
+    Any as Any_,
+)
 
-from pydantic import BaseModel as BaseModel_, Field, Extra
+from pydantic import (
+    BaseModel as BaseModel_,
+    ConfigDict,
+    Field,
+    SerializationInfo,
+    field_validator,
+    field_serializer,
+    ValidationError,
+)
+from pydantic.main import IncEx
+from pydantic_core import PydanticCustomError
+
+
+class AnyResource(BaseModel_):
+    model_config = ConfigDict(extra="allow")
+
+    resourceType: str
 
 
 
 class BaseModel(BaseModel_):
-    class Config:
-        extra = Extra.forbid
-        validate_assignment = True
-        allow_population_by_field_name = True
-
-    def dict(self, *args, **kwargs):
-        by_alias = kwargs.pop('by_alias', True)
-        return super().dict(*args, **kwargs, by_alias=by_alias)
+    model_config = ConfigDict(
+        # Extra attributes are disabled because fhir does not allow it
+        extra="forbid",
+        # Validation are applied while mutating the resource
+        validate_assignment=True,
+        # It's important for reserved keywords population in constructor (e.g. for_)
+        populate_by_name=True,
+        # Speed up initial load by lazy build
+        defer_build=True,
+        # It does not break anything, just for convinience
+        coerce_numbers_to_str=True,
+    )
+
+    def model_dump(
+        self,
+        *,
+        mode: Literal_["json", "python"] | str = "python",
+        include: IncEx = None,
+        exclude: IncEx = None,
+        context: Any_ | None = None,
+        by_alias: bool = True,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = True,
+        round_trip: bool = False,
+        warnings: bool | Literal_["none", "warn", "error"] = True,
+        serialize_as_any: bool = False,
+    ):
+        # Override default parameters for by_alias and exclude_none preserving function declaration
+        return super().model_dump(
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            context=context,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+            serialize_as_any=serialize_as_any,
+        )
+
+    @field_serializer("*")
+    @classmethod
+    def serialize_all_fields(cls, value: Any_, info: SerializationInfo):
+        if isinstance(value, list):
+            return [_serialize(v, info) for v in value]
+
+        return _serialize(value, info)
+
+    @field_validator("*")
+    @classmethod
+    def validate_all_fields(cls, value: Any_):
+        if isinstance(value, list):
+            return [_validate(v, index=index) for index, v in enumerate(value)]
+        return _validate(value)
+
+
+def _serialize(value: Any_, info: SerializationInfo):
+    # Custom serializer for AnyResource fields
+    kwargs = {
+        "mode": info.mode,
+        "include": info.include,
+        "exclude": info.exclude,
+        "context": info.context,
+        "by_alias": info.by_alias,
+        "exclude_unset": info.exclude_unset,
+        "exclude_defaults": info.exclude_defaults,
+        "exclude_none": info.exclude_none,
+        "round_trip": info.round_trip,
+        "serialize_as_any": info.serialize_as_any,
+    }
+    if isinstance(value, AnyResource):
+        return value.model_dump(**kwargs)  # type: ignore
+    if isinstance(value, BaseModel_):
+        return value.model_dump(**kwargs)  # type: ignore
+    return value
+
+
+def _validate(value: Any_, index: int | None = None):
+    # Custom validator for AnyResource fields
+    if isinstance(value, AnyResource):
+        try:
+            klass = globals()[value.resourceType]
+        except KeyError as exc:
+            raise ValidationError.from_exception_data(
+                "ImportError",
+                [
+                    {
+                        "loc": (index, "resourceType")
+                        if index is not None
+                        else ("resourceType",),
+                        "type": "value_error",
+                        "input": [value],
+                        "ctx": {"error": f"{value.resourceType} resource is not found"},
+                    }
+                ],
+            ) from exc
+
+        if not issubclass(klass, BaseModel) or "resourceType" not in klass.model_fields:
+            raise ValidationError.from_exception_data(
+                "ImportError",
+                [
+                    {
+                        "loc": (index, "resourceType")
+                        if index is not None
+                        else ("resourceType",),
+                        "type": "value_error",
+                        "input": [value],
+                        "ctx": {"error": f"{value.resourceType} is not a resource"},
+                    }
+                ],
+            )
+
+        try:
+            return klass(**value.model_dump())
+        except ValidationError as exc:
+            raise ValidationError.from_exception_data(
+                exc.title,
+                [
+                    {
+                        "loc": (index, *error["loc"])
+                        if index is not None
+                        else error["loc"],
+                        "type": error["type"],
+                        "input": error["input"],
+                        "ctx": error["ctx"]
+                    }
+                    for error in exc.errors()
+                ],
+            ) from exc
+
+    return value