Upgrade Pydantic to v2

.github/workflows/tests.yml

@@ -2,6 +2,8 @@ name: tests
 
 on:
   push:
+    branches:
+      - main
   pull_request:
     types: [opened, synchronize, reopened, edited]
 
@@ -11,14 +13,21 @@ jobs:
     if: github.repository == 'explosion/confection'
     strategy:
       fail-fast: false
+      # matrix:
+      #   os: [ubuntu-latest, windows-latest, macos-latest]
+      #   python_version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
       matrix:
-        os: [ubuntu-latest, windows-latest, macos-latest]
-        python_version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
         include:
-          - os: windows-2019
-            python_version: "3.6"
-          - os: ubuntu-20.04
-            python_version: "3.6"
+          - python_version: '3.7'
+            os: ubuntu-latest
+          - python_version: '3.8'
+            os: windows-latest
+          - python_version: '3.11'
+            os: ubuntu-latest
+          - python_version: '3.11'
+            os: windows-latest
+          - python_version: '3.11'
+            os: macos-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -39,8 +48,7 @@ jobs:
           python -m build --sdist
 
       - name: Run mypy
-        run: python -m mypy confection
-        if: matrix.python_version != '3.6'
+        run: python -m mypy confection/__init__.py confection/util.py
 
       - name: Delete source directory
         run: rm -rf confection
@@ -60,11 +68,25 @@ jobs:
       - name: Test import
         run: python -c "import confection" -Werror
 
-      - name: Install test requirements
-        run: python -m pip install -U -r requirements.txt
+      - name: Install test requirements with Pydantic v1
+        run: |
+          python -m pip install -U -r requirements.txt
+          python -m pip install -U pydantic==1.10.*
+
+      - name: Run tests for Pydantic v1
+        run: |
+          python -c "import pydantic; print(pydantic.VERSION)"
+          python -m pytest --pyargs confection
 
-      - name: Run tests
-        run: python -m pytest --pyargs confection -Werror
+      - name: Install test requirements with Pydantic v2
+        run: |
+          python -m pip install -U -r requirements.txt
+          python -m pip install -U pydantic
+
+      - name: Run tests for Pydantic v2
+        run: |
+          python -c "import pydantic; print(pydantic.VERSION)"
+          python -m pytest --pyargs confection
 
       - name: Test for import conflicts with hypothesis
         run: |

confection/__init__.py

@@ -1,15 +1,14 @@
 from typing import Union, Dict, Any, Optional, List, Tuple, Callable, Type, Mapping
-from typing import Iterable, Sequence, cast
+from typing import Iterable, Sequence, Set, TypeVar, TYPE_CHECKING, cast
 from types import GeneratorType
-from dataclasses import dataclass
+from dataclasses import dataclass, is_dataclass
 from configparser import ConfigParser, ExtendedInterpolation, MAX_INTERPOLATION_DEPTH
 from configparser import InterpolationMissingOptionError, InterpolationSyntaxError
 from configparser import NoSectionError, NoOptionError, InterpolationDepthError
 from configparser import ParsingError
 from pathlib import Path
-from pydantic import BaseModel, create_model, ValidationError, Extra
-from pydantic.main import ModelMetaclass
-from pydantic.fields import ModelField
+from pydantic import BaseModel, create_model, ValidationError
+from pydantic.fields import FieldInfo
 import srsly
 import catalogue
 import inspect
@@ -18,7 +17,7 @@
 import re
 import warnings
 
-from .util import Decorator, SimpleFrozenDict, SimpleFrozenList
+from .util import Decorator, SimpleFrozenDict, SimpleFrozenList, PYDANTIC_V2
 
 # Field used for positional arguments, e.g. [section.*.xyz]. The alias is
 # required for the schema (shouldn't clash with user-defined arg names)
@@ -666,19 +665,74 @@ def alias_generator(name: str) -> str:
     return name
 
 
-def copy_model_field(field: ModelField, type_: Any) -> ModelField:
+def copy_model_field(field: FieldInfo, type_: Any) -> FieldInfo:
     """Copy a model field and assign a new type, e.g. to accept an Any type
     even though the original value is typed differently.
     """
-    return ModelField(
-        name=field.name,
-        type_=type_,
-        class_validators=field.class_validators,
-        model_config=field.model_config,
-        default=field.default,
-        default_factory=field.default_factory,
-        required=field.required,
-    )
+    field_info = copy.deepcopy(field)
+    if PYDANTIC_V2:
+        field_info.annotation = type_
+    else:
+        field_info.type_ = type_  # type: ignore
+    return field_info
+
+
+def get_model_config_extra(model: Type[BaseModel]) -> str:
+    if PYDANTIC_V2:
+        extra = str(model.model_config.get("extra", "forbid"))
+    else:
+        extra = str(model.Config.extra) or "forbid"
+    assert isinstance(extra, str)
+    return extra
+
+
+
+_ModelT = TypeVar("_ModelT", bound=BaseModel)
+
+
+def model_validate(Schema: Type[_ModelT], data: Dict[str, Any]) -> _ModelT:
+    return Schema.model_validate(data) if PYDANTIC_V2 else Schema(**data)
+
+
+def model_construct(Schema: Type[_ModelT], fields_set: Optional[Set[str]], data: Dict[str, Any]) -> _ModelT:
+    return Schema.model_construct(fields_set, **data) if PYDANTIC_V2 else Schema.construct(fields_set, **data)
+
+
+def model_dump(instance: BaseModel) -> Dict[str, Any]:
+    return instance.model_dump() if PYDANTIC_V2 else instance.dict()
+
+
+def get_field_annotation(field: FieldInfo) -> Type:
+    return field.annotation if PYDANTIC_V2 else field.type_  # type: ignore
+
+
+def get_model_fields(Schema: Union[Type[BaseModel], BaseModel]) -> Dict[str, FieldInfo]:
+    return Schema.model_fields if PYDANTIC_V2 else Schema.__fields__  # type: ignore
+
+
+def get_model_fields_set(Schema: Union[Type[BaseModel], BaseModel]) -> Set[str]:
+    return Schema.model_fields_set if PYDANTIC_V2 else Schema.__fields_set__  # type: ignore
+
+
+def get_model_extra(model: BaseModel) -> Dict[str, FieldInfo]:
+    return model.model_extra or {} if PYDANTIC_V2 else {}
+
+
+def set_model_field(Schema: Type[BaseModel], key: str, field: FieldInfo):
+    if PYDANTIC_V2:
+        Schema.model_fields[key] = field
+    else:
+        Schema.__fields__[key] = field  # type: ignore
+
+
+def update_from_model_extra(shallow_result_dict: Dict[str, Any], result: BaseModel) -> None:
+    if PYDANTIC_V2:
+        if result.model_extra is not None:
+            shallow_result_dict.update(result.model_extra)
+
+
+def _safe_is_subclass(cls: type, expected: type) -> bool:
+    return inspect.isclass(cls) and issubclass(cls, BaseModel)
 
 
 class EmptySchema(BaseModel):
@@ -806,6 +860,7 @@ def _fill(
         resolve: bool = True,
         parent: str = "",
         overrides: Dict[str, Dict[str, Any]] = {},
+        resolved_object_keys: Set[str] = set()
     ) -> Tuple[
         Union[Dict[str, Any], Config], Union[Dict[str, Any], Config], Dict[str, Any]
     ]:
@@ -827,12 +882,14 @@ def _fill(
                 value = overrides[key_parent]
                 config[key] = value
             if cls.is_promise(value):
-                if key in schema.__fields__ and not resolve:
+                model_fields = get_model_fields(schema)
+                if key in model_fields and not resolve:
                     # If we're not resolving the config, make sure that the field
                     # expecting the promise is typed Any so it doesn't fail
                     # validation if it doesn't receive the function return value
-                    field = schema.__fields__[key]
-                    schema.__fields__[key] = copy_model_field(field, Any)
+                    field = model_fields[key]
+                    new_field = copy_model_field(field, Any)
+                    set_model_field(schema, key, new_field)
                 promise_schema = cls.make_promise_schema(value, resolve=resolve)
                 filled[key], validation[v_key], final[key] = cls._fill(
                     value,
@@ -841,6 +898,7 @@ def _fill(
                     resolve=resolve,
                     parent=key_parent,
                     overrides=overrides,
+                    resolved_object_keys=resolved_object_keys
                 )
                 reg_name, func_name = cls.get_constructor(final[key])
                 args, kwargs = cls.parse_args(final[key])
@@ -852,6 +910,9 @@ def _fill(
                     # We don't want to try/except this and raise our own error
                     # here, because we want the traceback if the function fails.
                     getter_result = getter(*args, **kwargs)
+
+                    if isinstance(getter_result, BaseModel) or is_dataclass(getter_result):
+                        resolved_object_keys.add(key)
                 else:
                     # We're not resolving and calling the function, so replace
                     # the getter_result with a Promise class
@@ -867,12 +928,12 @@ def _fill(
                     validation[v_key] = []
             elif hasattr(value, "items"):
                 field_type = EmptySchema
-                if key in schema.__fields__:
-                    field = schema.__fields__[key]
-                    field_type = field.type_
-                    if not isinstance(field.type_, ModelMetaclass):
-                        # If we don't have a pydantic schema and just a type
-                        field_type = EmptySchema
+                fields = get_model_fields(schema)
+                if key in fields:
+                    field = fields[key]
+                    annotation = get_field_annotation(field)
+                    if annotation is not None and _safe_is_subclass(annotation, BaseModel):
+                        field_type = annotation
                 filled[key], validation[v_key], final[key] = cls._fill(
                     value,
                     field_type,
@@ -898,21 +959,38 @@ def _fill(
         exclude = []
         if validate:
             try:
-                result = schema.parse_obj(validation)
+                result = model_validate(schema, validation)
             except ValidationError as e:
+
                 raise ConfigValidationError(
                     config=config, errors=e.errors(), parent=parent
                 ) from None
         else:
-            # Same as parse_obj, but without validation
-            result = schema.construct(**validation)
+            # Same as model_validate, but without validation
+            fields_set = set(get_model_fields(schema).keys())
+            result = model_construct(schema, fields_set, validation)
             # If our schema doesn't allow extra values, we need to filter them
             # manually because .construct doesn't parse anything
-            if schema.Config.extra in (Extra.forbid, Extra.ignore):
-                fields = schema.__fields__.keys()
-                exclude = [k for k in result.__fields_set__ if k not in fields]
+            if get_model_config_extra(schema) in ("forbid", "extra"):
+                result_field_names = get_model_fields_set(result)
+                exclude = [k for k in dict(result).keys() if k not in result_field_names]
         exclude_validation = set([ARGS_FIELD_ALIAS, *RESERVED_FIELDS.keys()])
-        validation.update(result.dict(exclude=exclude_validation))
+        # Do a shallow serialization first
+        # If any of the sub-objects are Pydantic models, first check if they
+        # were resolved earlier from a registry. If they weren't resolved
+        # they are part of a nested schema and need to be serialized with
+        # model.dict()
+        # Allows for returning Pydantic models from a registered function
+        shallow_result_dict = dict(result)
+        update_from_model_extra(shallow_result_dict, result)
+        result_dict = {}
+        for k, v in shallow_result_dict.items():
+            if k in exclude_validation:
+                continue
+            result_dict[k] = v
+            if isinstance(v, BaseModel) and k not in resolved_object_keys:
+                result_dict[k] = model_dump(v)
+        validation.update(result_dict)
         filled, final = cls._update_from_parsed(validation, filled, final)
         if exclude:
             filled = {k: v for k, v in filled.items() if k not in exclude}
@@ -946,6 +1024,8 @@ def _update_from_parsed(
             # Check numpy first, just in case. Use stringified type so that numpy dependency can be ditched.
             elif str(type(value)) == "<class 'numpy.ndarray'>":
                 final[key] = value
+            elif isinstance(value, BaseModel) and isinstance(final[key], BaseModel):
+                final[key] = value
             elif (
                 value != final[key] or not isinstance(type(value), type(final[key]))
             ) and not isinstance(final[key], GeneratorType):