DM-40002: Support pydantic v2

.github/workflows/build.yaml

@@ -82,7 +82,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Install dependencies
         run: |

.github/workflows/build_docs.yaml

@@ -18,7 +18,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: "pip"
           cache-dependency-path: "setup.cfg"
 

python/lsst/pipe/base/_task_metadata.py

@@ -27,11 +27,8 @@
 from collections.abc import Collection, Iterator, Mapping, Sequence
 from typing import Any, Protocol
 
-try:
-    from pydantic.v1 import BaseModel, Field, StrictBool, StrictFloat, StrictInt, StrictStr
-except ModuleNotFoundError:
-    from pydantic import BaseModel, Field, StrictBool, StrictFloat, StrictInt, StrictStr  # type: ignore
-
+from lsst.daf.butler._compat import _BaseModelCompat
+from pydantic import Field, StrictBool, StrictFloat, StrictInt, StrictStr
 
 _DEPRECATION_REASON = "Will be removed after v25."
 _DEPRECATION_VERSION = "v24"
@@ -59,7 +56,7 @@ def _isListLike(v: Any) -> bool:
     return isinstance(v, Sequence) and not isinstance(v, str)
 
 
-class TaskMetadata(BaseModel):
+class TaskMetadata(_BaseModelCompat):
     """Dict-like object for storing task metadata.
 
     Metadata can be stored at two levels: single task or task plus subtasks.
@@ -417,7 +414,7 @@ def get(self, key: str, default: Any = None) -> Any:
         key : `str`
             The key to retrieve. Can be dot-separated hierarchical.
         default
-            The value to return if the key doesnot exist.
+            The value to return if the key does not exist.
 
         Returns
         -------
@@ -574,4 +571,4 @@ def _validate_value(self, value: Any) -> tuple[str, Any]:
 
 
 # Needed because a TaskMetadata can contain a TaskMetadata.
-TaskMetadata.update_forward_refs()
+TaskMetadata.model_rebuild()

python/lsst/pipe/base/graph/quantumNode.py

@@ -35,13 +35,9 @@
     Quantum,
     SerializedQuantum,
 )
+from lsst.daf.butler._compat import _BaseModelCompat
 from lsst.utils.introspection import find_outside_stacklevel
 
-try:
-    from pydantic.v1 import BaseModel
-except ModuleNotFoundError:
-    from pydantic import BaseModel  # type: ignore
-
 from ..pipeline import TaskDef
 
 BuildId = NewType("BuildId", str)
@@ -145,7 +141,7 @@ def from_simple(
     ) -> QuantumNode:
         if recontitutedDimensions is not None:
             warnings.warn(
-                "The recontitutedDimensions argument is now ignored and may be removed after v 27",
+                "The recontitutedDimensions argument is now ignored and may be removed after v26",
                 category=FutureWarning,
                 stacklevel=find_outside_stacklevel("lsst.pipe.base"),
             )
@@ -159,7 +155,7 @@ def from_simple(
 _fields_set = {"quantum", "taskLabel", "nodeId"}
 
 
-class SerializedQuantumNode(BaseModel):
+class SerializedQuantumNode(_BaseModelCompat):
     """Model representing a `QuantumNode` in serializable form."""
 
     quantum: SerializedQuantum
@@ -168,10 +164,11 @@ class SerializedQuantumNode(BaseModel):
 
     @classmethod
     def direct(cls, *, quantum: dict[str, Any], taskLabel: str, nodeId: str) -> SerializedQuantumNode:
-        node = SerializedQuantumNode.__new__(cls)
-        setter = object.__setattr__
-        setter(node, "quantum", SerializedQuantum.direct(**quantum))
-        setter(node, "taskLabel", taskLabel)
-        setter(node, "nodeId", uuid.UUID(nodeId))
-        setter(node, "__fields_set__", _fields_set)
+        node = cls.model_construct(
+            __fields_set=_fields_set,
+            quantum=SerializedQuantum.direct(**quantum),
+            taskLabel=taskLabel,
+            nodeId=uuid.UUID(nodeId),
+        )
+
         return node

python/lsst/pipe/base/tests/mocks/_storage_class.py

@@ -34,10 +34,6 @@
 from collections.abc import Callable, Iterable, Mapping
 from typing import Any, cast
 
-try:
-    import pydantic.v1 as pydantic
-except ModuleNotFoundError:
-    import pydantic  # type: ignore
 from lsst.daf.butler import (
     DatasetComponent,
     Formatter,
@@ -50,6 +46,7 @@
     StorageClassDelegate,
     StorageClassFactory,
 )
+from lsst.daf.butler._compat import _BaseModelCompat
 from lsst.daf.butler.formatters.json import JsonFormatter
 from lsst.utils.introspection import get_full_type_name
 
@@ -82,7 +79,7 @@ def is_mock_name(name: str) -> bool:
 # access to complex real storage classes (and their pytypes) to test against.
 
 
-class MockDataset(pydantic.BaseModel):
+class MockDataset(_BaseModelCompat):
     """The in-memory dataset type used by `MockStorageClass`."""
 
     ref: SerializedDatasetRef
@@ -140,7 +137,7 @@ def make_derived(self, **kwargs: Any) -> MockDataset:
         `~lsst.daf.butler.SerializedDatasetType` to override in the result.
         """
         dataset_type_updates = {
-            k: kwargs.pop(k) for k in list(kwargs) if k in SerializedDatasetType.__fields__
+            k: kwargs.pop(k) for k in list(kwargs) if k in SerializedDatasetType.model_fields  # type: ignore
         }
         derived_dataset_type = self.dataset_type.copy(update=dataset_type_updates)
         derived_ref = self.ref.copy(update=dict(datasetType=derived_dataset_type))
@@ -156,7 +153,7 @@ def make_derived(self, **kwargs: Any) -> MockDataset:
         return self.copy(update=kwargs)
 
 
-class MockDatasetQuantum(pydantic.BaseModel):
+class MockDatasetQuantum(_BaseModelCompat):
     """Description of the quantum that produced a mock dataset."""
 
     task_label: str