From da0fdce95b7f0a95cec3e61271a988b12b668680 Mon Sep 17 00:00:00 2001
From: Vitaly Samigullin <vrs@pilosus.org>
Date: Fri, 22 Sep 2023 19:21:55 +0200
Subject: [PATCH] Support Pydantic v2 (#248)

---
 CHANGELOG.rst                    |  6 ++++++
 docs/code/flask_integration.py   | 12 ++++++++++--
 docs/code/pydantic_validation.py | 10 ++++++++--
 src/piny/__init__.py             |  7 ++++++-
 src/piny/validators.py           | 16 ++++++++++++++--
 tests/test_validators.py         | 13 +++++++++----
 6 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 94989905..508155a4 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,12 @@
 Changelog
 ---------
 
+v1.1.0 (2023-09-22)
+...................
+
+* Added: new validator `PydanticV2Validator` to support Pydantic v2
+
+
 v1.0.2 (2023-02-03)
 ...................
 
diff --git a/docs/code/flask_integration.py b/docs/code/flask_integration.py
index 0b2a8575..f63a8e2a 100644
--- a/docs/code/flask_integration.py
+++ b/docs/code/flask_integration.py
@@ -1,12 +1,20 @@
 from flask import Flask
 from flask.logging import default_handler
-from piny import YamlLoader, StrictMatcher, PydanticValidator
+from piny import YamlLoader, StrictMatcher, PydanticV2Validator
 from pydantic import BaseModel, validator
 from typing import Any, Dict, Optional
 from werkzeug.serving import run_simple
 import logging
 import sys
 
+# Watch out!
+# Pydantic V2 deprecated some model's methods:
+# https://docs.pydantic.dev/2.0/migration/
+#
+# For Pydantic v2 use `PydanticV2Validator`
+# For Pydantic v1 use `PydanticValidator`
+
+
 #
 # Validation
 #
@@ -78,7 +86,7 @@ def create_app(path: str) -> Flask:
     config = YamlLoader(
         path=path,
         matcher=StrictMatcher,
-        validator=PydanticValidator,
+        validator=PydanticV2Validator,
         schema=Configuration,
     ).load()
 
diff --git a/docs/code/pydantic_validation.py b/docs/code/pydantic_validation.py
index 91c3acc7..6d601331 100644
--- a/docs/code/pydantic_validation.py
+++ b/docs/code/pydantic_validation.py
@@ -1,6 +1,12 @@
 from pydantic import BaseModel
-from piny import PydanticValidator, StrictMatcher, YamlLoader
+from piny import PydanticV2Validator, StrictMatcher, YamlLoader
 
+# Watch out!
+# Pydantic V2 deprecated some model's methods:
+# https://docs.pydantic.dev/2.0/migration/
+#
+# For Pydantic v2 use `PydanticV2Validator`
+# For Pydantic v1 use `PydanticValidator`
 
 class DBModel(BaseModel):
     login: str
@@ -14,6 +20,6 @@ class ConfigModel(BaseModel):
 config = YamlLoader(
     path="database.yaml",
     matcher=StrictMatcher,
-    validator=PydanticValidator,
+    validator=PydanticV2Validator,
     schema=ConfigModel,
 ).load()
diff --git a/src/piny/__init__.py b/src/piny/__init__.py
index eba4799a..c2ef9554 100644
--- a/src/piny/__init__.py
+++ b/src/piny/__init__.py
@@ -1,4 +1,9 @@
 from .errors import ConfigError, LoadingError, ValidationError
 from .loaders import YamlLoader, YamlStreamLoader
 from .matchers import Matcher, MatcherWithDefaults, StrictMatcher
-from .validators import MarshmallowValidator, PydanticValidator, TrafaretValidator
+from .validators import (
+    MarshmallowValidator,
+    PydanticV2Validator,
+    PydanticValidator,
+    TrafaretValidator,
+)
diff --git a/src/piny/validators.py b/src/piny/validators.py
index 198d779a..babfc27e 100644
--- a/src/piny/validators.py
+++ b/src/piny/validators.py
@@ -25,9 +25,9 @@ def load(self, data: LoadedData, **params):
         pass  # pragma: no cover
 
 
-class PydanticValidator(Validator):
+class PydanticValidator(Validator):  # pragma: no cover
     """
-    Validator class for Pydantic library
+    Validator class for Pydantic Version 1
     """
 
     def load(self, data: LoadedData, **params):
@@ -37,6 +37,18 @@ def load(self, data: LoadedData, **params):
             raise ValidationError(origin=e, reason=str(e))
 
 
+class PydanticV2Validator(Validator):
+    """
+    Validator class for Pydantic Version 2
+    """
+
+    def load(self, data: LoadedData, **params):
+        try:
+            return self.schema(**data).model_dump()
+        except Exception as e:
+            raise ValidationError(origin=e, reason=str(e))
+
+
 class MarshmallowValidator(Validator):
     """
     Validator class for Marshmallow library
diff --git a/tests/test_validators.py b/tests/test_validators.py
index 1b5ee759..f16cde70 100644
--- a/tests/test_validators.py
+++ b/tests/test_validators.py
@@ -1,12 +1,14 @@
 from unittest import mock
 
+import pydantic
 import pytest
 import trafaret
 from marshmallow import Schema, fields
-from pydantic import BaseModel
+from packaging import version
 
 from piny import (
     MarshmallowValidator,
+    PydanticV2Validator,
     PydanticValidator,
     StrictMatcher,
     TrafaretValidator,
@@ -17,18 +19,21 @@
 
 from . import config_directory, config_map
 
+if version.parse(pydantic.__version__) >= version.parse("2.0"):
+    PydanticValidator = PydanticV2Validator
+
 #
 # Const
 #
 
 
-class PydanticDB(BaseModel):
+class PydanticDB(pydantic.BaseModel):
     host: str
     login: str
     password: str
 
 
-class PydanticConfig(BaseModel):
+class PydanticConfig(pydantic.BaseModel):
     db: PydanticDB
 
 
@@ -71,7 +76,7 @@ def test_pydantic_validator_success(name):
 
 @pytest.mark.parametrize("name", ["db"])
 def test_pydantic_validator_fail(name):
-    with pytest.raises(ValidationError, match=r"db -> password"):
+    with pytest.raises(ValidationError, match=r"db.*password"):
         YamlLoader(
             path=config_directory.joinpath("{}.yaml".format(name)),
             matcher=StrictMatcher,