Merge pull request #16 from fcurella/patch-15

Add support for Pydantic v2

.github/workflows/python-package.yml

@@ -19,7 +19,8 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        pydantic-version: ["<2", ">=2"]
 
     steps:
     - uses: actions/checkout@v3
@@ -39,6 +40,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         python -m pip install -e .
+        python -m pip install pydantic${{ matrix.pydantic-version }}
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
     - name: Test
       run: |
@@ -48,7 +50,7 @@ jobs:
       env:
         COVERALLS_PARALLEL: true
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        COVERALLS_FLAG_NAME: run-${{ matrix.python-version }}
+        COVERALLS_FLAG_NAME: run-${{ matrix.python-version }}-${{ matrix.pydantic-version }}
 
   finish:
     needs: [test]

Makefile

@@ -5,7 +5,7 @@ lint:
 	check-manifest
 
 test:
-	python -m pytest
+	python -m pytest  --pdb
 
 release:
 	check-manifest

datauri/__init__.py

@@ -4,7 +4,7 @@
 import textwrap
 from base64 import b64decode as decode64
 from base64 import b64encode as encode64
-from typing import Any, Dict, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, MutableMapping, Optional, Tuple, TypeVar, Union
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -160,15 +160,40 @@ def __get_validators__(cls):
         yield cls.validate
 
     @classmethod
-    def validate(cls, v: str) -> Self:
-        if not isinstance(v, str):
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: Any) -> Any:
+        from pydantic_core import core_schema
+
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
+    @classmethod
+    def validate(
+        cls,
+        value: str,
+        values: Optional[MutableMapping[str, Any]] = None,
+        config: Any = None,
+        field: Any = None,
+        **kwargs: Any,
+    ) -> Self:
+        if not isinstance(value, str):
             raise TypeError("string required")
 
-        m = cls(v)
+        m = cls(value)
         if not m.is_valid:
             raise ValueError("invalid data-uri format")
         return m
 
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, core_schema: MutableMapping[str, Any], handler: Any
+    ) -> Any:
+        core_schema.update(
+            pattern=DATA_URI_REGEX,
+            examples=[
+                "data:text/plain;charset=utf-8;base64,VGhlIHF1aWNrIGJyb3duIGZveCBqdW1wZWQgb3ZlciB0aGUgbGF6eSBkb2cu"
+            ],
+        )
+        return core_schema
+
     @classmethod
     def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
         # __modify_schema__ should mutate the dict it receives in place,

requirements.txt

@@ -1,6 +1,5 @@
 pytest
 coveralls
-pydantic<2.0
 flake8
 isort
 check-manifest

tests/test_pydantic.py

@@ -5,6 +5,7 @@
 pydantic = pytest.importorskip("pydantic")
 
 
+@pytest.mark.skipif(pydantic.__version__.rsplit(".", 3)[0] > "1", reason="pydantic v2")
 def test_pydantic():
     class Model(pydantic.BaseModel):
         content: DataURI
@@ -18,3 +19,21 @@ class Model(pydantic.BaseModel):
         'aGUgbGF6eSBkb2cu"}'
     )
     assert instance.dict() == {"content": DataURI(t)}
+
+
+@pytest.mark.skipif(
+    pydantic.__version__.rsplit(".", 3)[0] < "2", reason="pydantic v2 required"
+)
+def test_pydantic_v2():
+    class Model(pydantic.BaseModel):
+        content: DataURI
+
+    t = "data:text/plain;charset=utf-8;base64,VGhlIHF1aWNrIGJyb3duIGZveCBqdW1wZWQgb3ZlciB0aGUgbGF6eSBkb2cu"
+    instance = Model(content=t)
+    assert isinstance(instance.content, DataURI)
+    assert (
+        instance.model_dump_json()
+        == '{"content":"data:text/plain;charset=utf-8;base64,VGhlIHF1aWNrIGJyb3duIGZveCBqdW1wZWQgb3ZlciB0'
+        'aGUgbGF6eSBkb2cu"}'
+    )
+    assert instance.model_dump() == {"content": DataURI(t)}