Updated codebase to pydantic>=2.0

Fixes #63

Signed-off-by: Pedro Algarvio <[email protected]>

Author: s0undt3ch

.pre-commit-config.yaml

@@ -112,7 +112,7 @@ repos:
         additional_dependencies:
           - types-attrs
           - types-setuptools
-          - pydantic>=1.8.2,<2.0
+          - pydantic>=2.0,<3.0
           - types-aiofiles
           - aiostream
           - typing-extensions; python_version < "3.9"

changelog/63.improvement.rst

@@ -0,0 +1 @@
+Updated codebase to `pydantic>=2.0,<3.0`

examples/tests/pipelines/test_mnist_network.py

@@ -69,7 +69,8 @@ def test_pipeline(analytics_events_dump_directory: pathlib.Path):
         pytest.fail(f"Failed to find dumped events in {analytics_events_dump_directory}")
 
     contents = [
-        CollectedEvent.parse_obj(json.loads(i)) for i in dumpfile.read_text().strip().split("\n")
+        CollectedEvent.model_validate(json.loads(i))
+        for i in dumpfile.read_text().strip().split("\n")
     ]
     for event in contents:
         assert "evaluation" in event.data

examples/tests/pipelines/test_mnist_notebook.py

@@ -89,7 +89,8 @@ def test_pipeline(analytics_events_dump_directory: pathlib.Path):
         pytest.fail(f"Failed to find dumped events in {analytics_events_dump_directory}")
 
     contents = [
-        CollectedEvent.parse_obj(json.loads(i)) for i in dumpfile.read_text().strip().split("\n")
+        CollectedEvent.model_validate(json.loads(i))
+        for i in dumpfile.read_text().strip().split("\n")
     ]
     for event in contents:
         assert "evaluation" in event.data

requirements/base.txt

@@ -5,6 +5,6 @@ backoff
 drain3
 importlib-metadata>=3.4.0,<5.0.0; python_version < "3.8"
 psutil
-pydantic>=1.8.2,<2.0
+pydantic>=2.0,<3.0
 salt>=3005
 typing-extensions; python_version < "3.9.2"

src/saf/collect/beacons.py

@@ -19,7 +19,7 @@
 from typing import TypeVar
 from typing import Union
 
-from pydantic import validator
+from pydantic import field_validator
 
 from saf.models import CollectConfigBase
 from saf.models import CollectedEvent
@@ -60,7 +60,7 @@ def _convert_stamp(stamp: str) -> datetime:
             _stamp = datetime.strptime(stamp, "%Y-%m-%dT%H:%M:%S.%f").replace(tzinfo=timezone.utc)
         return _stamp
 
-    @validator("stamp")
+    @field_validator("stamp")
     @classmethod
     def _validate_stamp(cls: Type[BCE], value: Union[str, datetime]) -> datetime:
         if isinstance(value, datetime):

src/saf/collect/file.py

@@ -22,7 +22,7 @@
 from saf.models import CollectedEvent
 from saf.models import PipelineRunContext
 
-if sys.version_info < (3, 9, 2):
+if sys.version_info < (3, 12):
     from typing_extensions import TypedDict
 else:
     from typing import TypedDict  # type: ignore[attr-defined,no-redef]

src/saf/forward/disk.py

@@ -56,10 +56,10 @@ async def forward(
         dest = config.path / config.filename
         dest.touch()
         async with aiofiles.open(dest, "a", encoding="utf-8") as wfh:
-            wrote = await wfh.write(f"{event.json(indent=indent)}\n")
+            wrote = await wfh.write(f"{event.model_dump_json(indent=indent)}\n")
     else:
         file_count = len(list(config.path.iterdir()))
         dest = config.path / f"event-dump-{file_count + 1}.json"
         async with aiofiles.open(dest, "w", encoding="utf-8") as wfh:
-            wrote = await wfh.write(event.json(indent=indent))
+            wrote = await wfh.write(event.model_dump_json(indent=indent))
     log.debug("Wrote %s bytes to %s", wrote, dest)

src/saf/forward/test.py

@@ -18,7 +18,7 @@
 from typing import Optional
 from typing import Type
 
-from pydantic import root_validator
+from pydantic import model_validator
 
 from saf.models import CollectedEvent
 from saf.models import ForwardConfigBase
@@ -38,7 +38,7 @@ class TestForwardConfig(ForwardConfigBase):
     dump_event: bool = False
     add_event_to_shared_cache: bool = False
 
-    @root_validator
+    @model_validator(mode="before")
     @classmethod
     def _check_mutually_exclusive_parameters(
         cls: Type[TestForwardConfig], values: dict[str, Any]
@@ -81,7 +81,7 @@ async def forward(
     elif config.path:
         if config.message:
             if config.dump_event:
-                dump_text = json.dumps({config.message: event.dict()})
+                dump_text = json.dumps({config.message: event.model_dump()})
             else:
                 dump_text = config.message
         elif config.dump_event: