Add support for classvars (#292)

Author: zmievsa

CHANGELOG.md

@@ -5,6 +5,12 @@ Please follow [the Keep a Changelog standard](https://keepachangelog.com/en/1.0.
 
 ## [Unreleased]
 
+## [5.4.4]
+
+### Fixed
+
+* Fixed KeyError when versioning Pydantic models containing ClassVar annotations
+
 ## [5.4.3]
 
 ### Fixed

cadwyn/schema_generation.py

@@ -14,6 +14,7 @@
 from typing import (
     TYPE_CHECKING,
     Annotated,
+    ClassVar,
     Generic,
     Union,
     _BaseGenericAlias,  # pyright: ignore[reportAttributeAccessIssue]
@@ -41,6 +42,7 @@
 from pydantic._internal._known_annotated_metadata import collect_known_metadata
 from pydantic._internal._typing_extra import try_eval_type as pydantic_try_eval_type
 from pydantic.fields import ComputedFieldInfo, FieldInfo
+from pydantic_core import PydanticUndefined
 from typing_extensions import (
     Any,
     Doc,
@@ -316,7 +318,7 @@ def _rebuild_annotated(name: str):
             field_name,
         )
         for field_name in model.__annotations__
-        if field_name in defined_fields
+        if field_name in defined_fields and field_name in model.model_fields
     }
 
     main_attributes = fields | validators
@@ -592,7 +594,12 @@ def _change_version_of_a_non_container_annotation(self, annotation: Any) -> Any:
         from typing_inspection.typing_objects import is_any, is_newtype, is_typealiastype
 
         if isinstance(annotation, (types.GenericAlias, _BaseGenericAlias)):
-            return get_origin(annotation)[tuple(self.change_version_of_annotation(arg) for arg in get_args(annotation))]
+            origin = get_origin(annotation)
+            args = get_args(annotation)
+            # Classvar does not support generic tuple arguments
+            if origin is ClassVar:
+                return ClassVar[self.change_version_of_annotation(args[0])]
+            return origin[tuple(self.change_version_of_annotation(arg) for arg in get_args(annotation))]
         elif is_typealiastype(annotation):
             if (
                 annotation.__module__ is not None and (annotation.__module__.startswith("pydantic."))
@@ -948,11 +955,20 @@ def _add_field_to_model(
             f'in "{version_change_name}" but there is already a field with that name.',
         )
 
-    field = PydanticFieldWrapper(
-        alter_schema_instruction.field, alter_schema_instruction.field.annotation, alter_schema_instruction.name
-    )
-    model.fields[alter_schema_instruction.name] = field
-    model.annotations[alter_schema_instruction.name] = alter_schema_instruction.field.annotation
+    # Special handling for ClassVar fields
+    if get_origin(alter_schema_instruction.field.annotation) is ClassVar:
+        # ClassVar fields should not be in model.fields, only in annotations and other_attributes
+        model.annotations[alter_schema_instruction.name] = alter_schema_instruction.field.annotation
+        # Set the actual ClassVar value in other_attributes
+        if alter_schema_instruction.field.default is not PydanticUndefined:
+            model.other_attributes[alter_schema_instruction.name] = alter_schema_instruction.field.default
+    else:
+        # Regular field handling
+        field = PydanticFieldWrapper(
+            alter_schema_instruction.field, alter_schema_instruction.field.annotation, alter_schema_instruction.name
+        )
+        model.fields[alter_schema_instruction.name] = field
+        model.annotations[alter_schema_instruction.name] = alter_schema_instruction.field.annotation
 
 
 def _change_field_in_model(
@@ -1085,6 +1101,11 @@ def _delete_field_from_model(model: _PydanticModelWrapper, field_name: str, vers
         validator = model.validators[field_name]
         model.validators[field_name].is_deleted = True
         model.annotations.pop(field_name, None)
+    elif field_name in model.annotations and get_origin(model.annotations[field_name]) is ClassVar:
+        # Handle ClassVar fields - they exist in annotations but not in model.fields
+        model.annotations.pop(field_name)
+        # Also remove the attribute from other_attributes if it exists there
+        model.other_attributes.pop(field_name, None)
     else:
         raise InvalidGenerationInstructionError(
             f'You tried to delete a field "{field_name}" from "{model.name}" '

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cadwyn"
-version = "5.4.3"
+version = "5.4.4"
 description = "Production-ready community-driven modern Stripe-like API versioning in FastAPI"
 authors = [{ name = "Stanislav Zmiev", email = "zmievsa@gmail.com" }]
 license = "MIT"

tests/test_schema_generation/test_schema_field.py

@@ -1,6 +1,6 @@
 import re
 from enum import Enum, auto
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, ClassVar, Literal, Union, get_origin
 
 import pytest
 from pydantic import BaseModel, Field, StringConstraints, ValidationError, computed_field, conint, constr
@@ -807,3 +807,81 @@ def image_url(self) -> str:
 
     assert not hasattr(old_instance, "image_url")
     assert "image_url" not in old_instance.model_dump()
+
+
+def test__schema_with_classvar__should_be_recreated_in_older_version(create_runtime_schemas: CreateRuntimeSchemas):
+    class SchemaWithClassVar(BaseModel):
+        regular_field: str
+        class_level_config: ClassVar[str] = "default_config"
+
+    schemas = create_runtime_schemas(version_change())
+
+    latest_model = schemas["2001-01-01"][SchemaWithClassVar]
+    old_model = schemas["2000-01-01"][SchemaWithClassVar]
+
+    latest_instance = latest_model(regular_field="test")
+    old_instance = old_model(regular_field="test")
+
+    assert latest_model.class_level_config == "default_config"
+    assert old_model.class_level_config == "default_config"
+
+    assert "class_level_config" not in latest_instance.model_dump()
+    assert "class_level_config" not in old_instance.model_dump()
+
+    assert latest_instance.regular_field == "test"
+    assert old_instance.regular_field == "test"
+
+    assert "class_level_config" in latest_model.__annotations__
+    assert "class_level_config" in old_model.__annotations__
+
+    assert get_origin(latest_model.__annotations__["class_level_config"]) is ClassVar
+    assert get_origin(old_model.__annotations__["class_level_config"]) is ClassVar
+
+
+def test__schema_with_classvar__remove_classvar(create_runtime_schemas: CreateRuntimeSchemas):
+    class SchemaWithClassVar(BaseModel):
+        regular_field: str
+        class_level_config: ClassVar[str] = "default_config"
+
+    schemas = create_runtime_schemas(version_change(schema(SchemaWithClassVar).field("class_level_config").didnt_exist))
+
+    latest_model = schemas["2001-01-01"][SchemaWithClassVar]
+    latest_instance = latest_model(regular_field="test")
+
+    assert latest_model.class_level_config == "default_config"
+    assert "class_level_config" not in latest_instance.model_dump()
+
+    old_model = schemas["2000-01-01"][SchemaWithClassVar]
+    old_instance = old_model(regular_field="test")
+
+    assert not hasattr(old_model, "class_level_config")
+    assert "class_level_config" not in old_instance.model_dump()
+
+    assert latest_instance.regular_field == "test"
+    assert old_instance.regular_field == "test"
+
+
+def test__schema_with_classvar__add_classvar_field(create_runtime_schemas: CreateRuntimeSchemas):
+    class SchemaWithoutClassVar(BaseModel):
+        regular_field: str
+
+    schemas = create_runtime_schemas(
+        version_change(),
+        version_change(
+            schema(SchemaWithoutClassVar)
+            .field("new_config")
+            .existed_as(type=ClassVar[str], info=Field(default="added_config")),
+            schema(SchemaWithoutClassVar).field("new_config_without_value").existed_as(type=ClassVar[str]),
+        ),
+    )
+    latest_model = schemas["2002-01-01"][SchemaWithoutClassVar]
+    mid_model = schemas["2001-01-01"][SchemaWithoutClassVar]
+    old_model = schemas["2000-01-01"][SchemaWithoutClassVar]
+
+    assert not hasattr(latest_model, "new_config")
+    assert mid_model.new_config == "added_config"
+    assert old_model.new_config == "added_config"
+
+    assert not hasattr(latest_model, "new_config_without_value")
+    assert not hasattr(mid_model, "new_config_without_value")
+    assert not hasattr(old_model, "new_config_without_value")