chore: address pydantic deprecated methods

canonical · Jul 1, 2023 · b3d0d49 · b3d0d49
1 parent 093ee33
commit b3d0d49
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/lib/charms/data_platform_libs/v0/data_models.py b/lib/charms/data_platform_libs/v0/data_models.py
@@ -231,7 +231,7 @@ def write(relation_data: RelationDataContent, model: BaseModel):
         relation_data: pointer to the relation databag
         model: instance of pydantic model to be written
     """
-    for key, value in model.dict(exclude_none=True).items():
+    for key, value in model.model_dump(exclude_none=True).items():
         relation_data[key.replace("_", "-")] = (
             str(value) if isinstance(value, str) or isinstance(value, int) else json.dumps(value)
         )
@@ -248,7 +248,7 @@ def read(relation_data: MutableMapping[str, str], obj: Type[T]) -> T:
         **{
             field_name: (
                 relation_data[parsed_key]
-                if field.type_ in [int, str, float]
+                if field.annotation in [int, str, float]
                 else json.loads(relation_data[parsed_key])
             )
             for field_name, field in obj.__fields__.items()  # pyright: ignore[reportGeneralTypeIssues]

diff --git a/lib/charms/data_platform_libs/v0/upgrade.py b/lib/charms/data_platform_libs/v0/upgrade.py
@@ -17,7 +17,7 @@
 import json
 import logging
 from abc import ABC, abstractmethod
-from typing import Optional, Union
+from typing import Optional
 
 from ops.charm import (
     ActionEvent,
@@ -28,7 +28,7 @@
 )
 from ops.framework import EventBase, Object
 from ops.model import Relation
-from pydantic import BaseModel, root_validator, validator
+from pydantic import BaseModel, field_validator, model_validator
 
 # The unique Charmhub library identifier, never change it
 LIBID = "156258aefb79435a93d933409a8c8684"
@@ -265,7 +265,7 @@ class KafkaDependenciesModel(BaseModel):
         model = KafkaDependenciesModel(**deps)
 
         # exporting back validated deps
-        print(model.dict())
+        print(model.model_dump())
         ```
     """
 
@@ -274,7 +274,7 @@ class KafkaDependenciesModel(BaseModel):
     upgrade_supported: str
     version: str
 
-    @validator("dependencies", "upgrade_supported")
+    @field_validator("dependencies", "upgrade_supported")
     @classmethod
     def dependencies_validator(cls, value):
         """Validates values with dependencies for multiple special characters."""
@@ -293,9 +293,9 @@ def dependencies_validator(cls, value):
 
         return value
 
-    @root_validator(skip_on_failure=True)
+    @model_validator(mode="before")
     @classmethod
-    def version_upgrade_supported_validator(cls, values) -> dict[str, Union[dict[str, str], str]]:
+    def version_upgrade_supported_validator(cls, values):
         """Validates specified `version` meets `upgrade_supported` requirement."""
         if not verify_requirements(
             version=values.get("version"), requirement=values.get("upgrade_supported")
@@ -426,7 +426,7 @@ def stored_dependencies(self) -> Optional[BaseModel]:
         if not (deps := self.peer_relation.data[self.charm.app].get("dependencies", "")):
             return None
 
-        return type(self.dependency_model)(**json.loads(deps))
+        return type(self.dependency_model).model_validate_json(deps)
 
     @property
     def upgrade_stack(self) -> Optional[list[int]]:
@@ -485,7 +485,7 @@ def _on_upgrade_created(self, _: RelationCreatedEvent) -> None:
 
         # persisting dependencies to the relation data
         self.peer_relation.data[self.charm.app].update(
-            {"dependencies": json.dumps(self.dependency_model.dict())}
+            {"dependencies": self.dependency_model.model_dump_json()}
         )
 
     def _on_pre_upgrade_check_action(self, event: ActionEvent):