Update for pydantic v2

DiamondLightSource · May 15, 2024 · b39faa5 · b39faa5
1 parent 3129c18
commit b39faa5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/src/hdf5_reader_service/model.py b/src/hdf5_reader_service/model.py
@@ -3,7 +3,6 @@
 
 import h5py as h5
 from pydantic import BaseModel
-from pydantic.generics import GenericModel
 
 
 class DatasetMacroStructure(BaseModel):
@@ -55,7 +54,7 @@ class ShapeMetadata(BaseModel):
 T = TypeVar("T")
 
 
-class ValidNode(GenericModel, Generic[T]):
+class ValidNode(BaseModel, Generic[T]):
     contents: T
     subnodes: List["DataTree"] = []
 
@@ -69,10 +68,10 @@ class InvalidNode(BaseModel):
     reason: InvalidNodeReason
 
 
-class DataTree(GenericModel, Generic[T]):
+class DataTree(BaseModel, Generic[T]):
     name: str
     valid: bool
     node: Union[InvalidNode, ValidNode[T]]
 
 
-ValidNode.update_forward_refs()
+ValidNode.model_rebuild()
diff --git a/src/hdf5_reader_service/utils.py b/src/hdf5_reader_service/utils.py
@@ -33,7 +33,7 @@ def default(content):
                 return content.decode("utf-8")
             elif isinstance(content, BaseModel):
                 # Handle the pydantic model case
-                return content.dict()
+                return content.model_dump()
         raise TypeError
 
     # Not all numpy dtypes are supported by orjson.