Refactoring to use actual Pydantic V2 API (#190)

Author: JoOkuma

Diff for: .github/workflows/docs.yml

@@ -38,7 +38,6 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install ".[docs]"
-          pip install "autodoc_pydantic==1.9.1"
       - name: Set environment
         run: |
           REPO_OWNER="${GITHUB_REPOSITORY%%/*}"

Diff for: pyproject.toml

@@ -33,7 +33,8 @@ dependencies = [
     "psycopg2-binary >=2.9.6",
     "psygnal >=0.9.0",
     "pyarrow >=16.1.0,<20",
-    "pydantic >= 2",
+    "pydantic >=2",
+    "pydantic-settings",
     "pydot >=2.0.0",
     "qtawesome >=1.3.1",
     "qtpy",
@@ -53,8 +54,7 @@ dependencies = [
 [project.optional-dependencies]
 docs = [
     "sphinx == 6.2.1",
-    "autodoc_pydantic",  # NOT WORKING, needs to install 'autodoc_pydantic==1.9.1' to build docs
-    "pydantic-settings",
+    "autodoc_pydantic",
     "furo",
     "myst-parser >= 2.0.0",
     "nbsphinx >= 0.9.3",
@@ -105,6 +105,7 @@ pillow = ">=10.0.0"
 psycopg2-binary = ">=2.9.6"
 psygnal = ">=0.9.0"
 pydantic = ">=2"
+pydantic-settings = "*"
 pydot = ">=2.0.0"
 qtawesome = ">=1.3.1"
 qtpy = "*"

Diff for: ultrack/api/_test/test_api.py

@@ -77,7 +77,7 @@ def experiment_instance(
 
     experiment = Experiment(
         name="PyTest",
-        config=config_instance,
+        config=config_instance.model_dump(by_alias=True),
         data_url=ome_zarr_dataset_path,
         image_channel_or_path="image",
         edges_channel_or_path="edges",
@@ -103,7 +103,7 @@ def test_config():
     default_config = MainConfig()
     default_config.data_config = None
 
-    assert response.json() == default_config.dict()
+    assert response.json() == default_config.model_dump(by_alias=True)
 
 
 def test_manual_segment(experiment_instance: Experiment):
@@ -130,7 +130,7 @@ def test_manual_segment(experiment_instance: Experiment):
     with client.websocket_connect("/segment/manual") as websocket:
         json_exp = json.loads(
             json.dumps(
-                experiment.dict(),
+                experiment.model_dump(),
                 default=lambda o: o.isoformat()
                 if isinstance(o, (datetime.date, datetime.datetime))
                 else None,
@@ -139,7 +139,7 @@ def test_manual_segment(experiment_instance: Experiment):
         websocket.send_json({"experiment": json_exp})
         while experiment.status != ExperimentStatus.SUCCESS:
             response = websocket.receive_json()
-            experiment = Experiment.parse_obj(response)
+            experiment = Experiment.model_validate(response)
             assert experiment.status != ExperimentStatus.ERROR
 
     tracks_df_api, graph_api = to_tracks_layer(experiment.get_config())
@@ -179,7 +179,7 @@ def test_auto_detect(
     with client.websocket_connect("/segment/auto_detect") as websocket:
         json_exp = json.loads(
             json.dumps(
-                experiment.dict(),
+                experiment.model_dump(),
                 default=lambda o: o.isoformat()
                 if isinstance(o, (datetime.date, datetime.datetime))
                 else None,
@@ -201,7 +201,7 @@ def test_auto_detect(
         websocket.send_json(json_request)
         while experiment.status != ExperimentStatus.SUCCESS:
             response = websocket.receive_json()
-            experiment = Experiment.parse_obj(response)
+            experiment = Experiment.model_validate(response)
             assert experiment.status != ExperimentStatus.ERROR
 
     detection = np.zeros_like(image_data, dtype=float)
@@ -250,7 +250,7 @@ def test_from_labels(experiment_instance: Experiment, label_to_edges_kwargs: dic
     with client.websocket_connect("/segment/labels") as websocket:
         json_exp = json.loads(
             json.dumps(
-                experiment.dict(),
+                experiment.model_dump(),
                 default=lambda o: o.isoformat()
                 if isinstance(o, (datetime.date, datetime.datetime))
                 else o,
@@ -265,7 +265,7 @@ def test_from_labels(experiment_instance: Experiment, label_to_edges_kwargs: dic
             label_to_edges_kwargs = {}
         while experiment.status != ExperimentStatus.SUCCESS:
             response = websocket.receive_json()
-            experiment = Experiment.parse_obj(response)
+            experiment = Experiment.model_validate(response)
             assert experiment.status != ExperimentStatus.ERROR
 
     # compare the results with the ones obtained from the ultrack module
@@ -293,7 +293,7 @@ def test_output_experiment(experiment_instance: Experiment):
     with client.websocket_connect("/segment/auto_detect") as websocket:
         json_exp = json.loads(
             json.dumps(
-                experiment_instance.dict(),
+                experiment_instance.model_dump(),
                 default=lambda o: o.isoformat()
                 if isinstance(o, (datetime.date, datetime.datetime))
                 else None,
@@ -306,7 +306,7 @@ def test_output_experiment(experiment_instance: Experiment):
 
         while experiment_instance.status != ExperimentStatus.SUCCESS:
             response = websocket.receive_json()
-            experiment_instance = Experiment.parse_obj(response)
+            experiment_instance = Experiment.model_validate(response)
             assert experiment_instance.status != ExperimentStatus.ERROR
 
         assert "robust_invert" in experiment_instance.err_log
@@ -353,7 +353,7 @@ def test_available_configs(experiment_instance: Experiment):
 
             while experiment.status != ExperimentStatus.SUCCESS:
                 response = websocket.receive_json()
-                experiment = Experiment.parse_obj(response)
+                experiment = Experiment.model_validate(response)
                 assert experiment.status != ExperimentStatus.ERROR
 
 
@@ -377,7 +377,7 @@ def test_available_configs(experiment_instance: Experiment):
 #
 #         while experiment_instance.status != ExperimentStatus.ERROR:
 #             response = websocket.receive_json()
-#             experiment_instance = Experiment.parse_obj(response)
+#             experiment_instance = Experiment.model_validate(response)
 #
 #         assert "Exception" in experiment_instance.err_log
 #         assert "Traceback" in experiment_instance.err_log
@@ -399,9 +399,9 @@ def test_available_configs(experiment_instance: Experiment):
 #         )
 #         websocket.send_json({"experiment": json_exp})
 #         response = websocket.receive_json()
-#         experiment_instance = Experiment.parse_obj(response)
+#         experiment_instance = Experiment.model_validate(response)
 #         response = client.post(f"/stop/{experiment_instance.id}")
 #         assert response.status_code == 200
 #         response = websocket.receive_json()
-#         experiment_instance = Experiment.parse_obj(response)
+#         experiment_instance = Experiment.model_validate(response)
 #         assert experiment_instance.status == ExperimentStatus.ERROR

Diff for: ultrack/api/app.py

@@ -151,7 +151,7 @@ async def finish_experiment(ws: WebSocket, exp: Experiment) -> Experiment:
         The experiment instance to be finished.
 
     """
-    exp.end_time = datetime.now()
+    exp.end_time = datetime.now().isoformat()
     exp.status = ExperimentStatus.SUCCESS
     update_experiment(exp)
     app.state.queue.task_done()
@@ -184,7 +184,7 @@ async def get_default_config() -> Dict:
     """
     config = MainConfig()
     config.data_config = None
-    return config.dict()
+    return config.model_dump(by_alias=True)
 
 
 @app.get("/config/available")
@@ -201,7 +201,7 @@ async def get_available_configs() -> Dict:
 
     experiment = {
         "name": "Unnamed Experiment",
-        "config": default_config.dict(),
+        "config": default_config.model_dump(by_alias=True),
     }
 
     auto_detect_config = {
@@ -316,7 +316,7 @@ async def auto_detect(websocket: WebSocket) -> None:
         )
         return
 
-    experiment = Experiment.parse_obj(data["experiment"])
+    experiment = Experiment.model_validate(data["experiment"])
 
     try:
         detect_foreground_kwargs = data["detect_foreground_kwargs"]
@@ -432,7 +432,7 @@ async def manual_segment(websocket: WebSocket) -> None:
         experiment = json.loads(data["experiment"])
     else:
         experiment = data["experiment"]
-    experiment = Experiment.parse_obj(experiment)
+    experiment = Experiment.model_validate(experiment)
 
     async with UltrackWebsocketLogger(websocket, experiment):
         await start_experiment(websocket, experiment)
@@ -495,7 +495,7 @@ async def auto_from_labels(websocket: WebSocket) -> None:
         )
         return
 
-    experiment = Experiment.parse_obj(data["experiment"])
+    experiment = Experiment.model_validate(data["experiment"])
 
     try:
         label_to_edges_kwargs = data["label_to_edges_kwargs"]

Diff for: ultrack/api/database.py

@@ -6,7 +6,7 @@
 from typing import Optional
 
 import sqlalchemy as sqla
-from pydantic.v1 import BaseModel, Json, validator
+from pydantic import BaseModel, Json, ValidationInfo, field_validator
 from sqlalchemy import JSON, Column, Enum, Integer, String, Text
 from sqlalchemy.orm import declarative_base, sessionmaker
 
@@ -22,7 +22,7 @@
 def _clean_db_on_exit():
     try:
         Session._temp_dir.cleanup()
-    except:
+    except Exception:
         pass
 
 
@@ -104,9 +104,9 @@ class Experiment(BaseModel):
         The experiment status. Defaults to ExperimentStatus.NOT_PERSISTED.
     name : str
         The experiment name. Defaults to "Untitled Experiment".
-    start_time : Optional[datetime]
+    start_time : Optional[str]
         The experiment start time. Defaults to the datetime where it was created.
-    end_time : Optional[datetime]
+    end_time : Optional[str]
         The experiment end time. Defaults to None and is set when the experiment
         finishes.
     std_log : str
@@ -153,34 +153,37 @@ class Experiment(BaseModel):
     tracks: Optional[Json] = None
 
     def get_config(self) -> MainConfig:
-        config = MainConfig.parse_obj(self.config)
+        config = MainConfig.model_validate(self.config)
         config.data_config = settings.ultrack_data_config
         return config
 
-    @validator("status", pre=True, always=True)
-    def check_if_id_is_valid(cls, v, values, **kwargs):
-        if (
-            v == ExperimentStatus.NOT_PERSISTED
-            and "id" in values
-            and values["id"] is not None
-        ):
+    @field_validator("status", mode="before")
+    @classmethod
+    def check_if_id_is_valid(cls, v, info: ValidationInfo) -> ExperimentStatus:
+        exp_id = info.data.get("id")
+        if v == ExperimentStatus.NOT_PERSISTED and exp_id is not None:
             raise ValueError(
                 "The id cannot be set if the experiment was never "
                 "persisted in the database."
             )
-        elif v != ExperimentStatus.NOT_PERSISTED and "id" not in values:
+        elif v != ExperimentStatus.NOT_PERSISTED and exp_id is None:
             raise ValueError(
                 "The id must be set if the experiment was persisted in the database."
             )
         return v
 
-    @validator("start_time", pre=True, always=True)
-    def set_start_time(cls, v):
+    @field_validator("start_time", mode="before")
+    @classmethod
+    def set_start_time(cls, v) -> str:
         return v or datetime.now().isoformat()
 
-    @validator("end_time", always=True)
-    def set_end_time(cls, v, values):
-        if values["status"] in [ExperimentStatus.SUCCESS, ExperimentStatus.FAILED]:
+    @field_validator("end_time")
+    @classmethod
+    def set_end_time(cls, v, info: ValidationInfo) -> Optional[str]:
+        if info.data.get("status") in [
+            ExperimentStatus.SUCCESS,
+            ExperimentStatus.FAILED,
+        ]:
             return v or datetime.now().isoformat()
         return None
 
@@ -262,7 +265,7 @@ def update_experiment(experiment: Experiment) -> None:
     experiment_db = session.query(ExperimentDB).filter_by(id=experiment.id).first()
     if experiment_db is None:
         raise ValueError(f"Experiment {experiment.id} not found.")
-    for key, value in experiment.dict().items():
+    for key, value in experiment.model_dump().items():
         if key != "id":
             setattr(experiment_db, key, value)
     session.commit()

Diff for: ultrack/api/settings.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Union
 
-from pydantic.v1 import BaseSettings
+from pydantic_settings import BaseSettings
 
 from ultrack.config import DataConfig
 

Diff for: ultrack/cli/config.py

@@ -16,4 +16,4 @@ def config_cli(output_path: Path) -> None:
         raise ValueError(f"{output_path} already exists.")
 
     with open(output_path, mode="w") as f:
-        toml.dump(config.dict(by_alias=True), f)
+        toml.dump(config.model_dump(by_alias=True), f)

Diff for: ultrack/config/_test/test_config.py

@@ -3,7 +3,6 @@
 
 import pytest
 import toml
-from pydantic.v1 import ValidationError
 
 from ultrack.config import load_config
 
@@ -14,7 +13,7 @@ def _assert_input_in_target(input: Dict, target: Dict) -> None:
         if isinstance(input[k], dict):
             _assert_input_in_target(input[k], target[k])
         else:
-            assert target[k] == input[k]
+            assert target[k] == input[k], f"Key {k} not equal"
 
 
 def _format_config(config: Dict) -> None:
@@ -29,7 +28,7 @@ def test_config_content(config_path: Path, config_content: Dict[str, Any]) -> No
     """Tests if content is loaded correctly"""
     config = load_config(config_path)
     _format_config(config_content)
-    _assert_input_in_target(config_content, config.dict())
+    _assert_input_in_target(config_content, config.model_dump())
 
 
 def test_invalid_config_content(tmp_path: Path, config_content: Dict[str, Any]) -> None:
@@ -40,5 +39,5 @@ def test_invalid_config_content(tmp_path: Path, config_content: Dict[str, Any])
     with open(path, mode="w") as f:
         toml.dump(config_content, f)
 
-    with pytest.raises(ValidationError):
+    with pytest.raises(ValueError):
         load_config(path)

Diff for: ultrack/config/config.py

@@ -3,7 +3,7 @@
 from typing import Optional, Union
 
 import toml
-from pydantic.v1 import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from ultrack.config.dataconfig import DataConfig
 from ultrack.config.segmentationconfig import SegmentationConfig
@@ -38,13 +38,13 @@ class LinkingConfig(BaseModel):
     the segmentation masks of neighboring segments
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(extra="forbid")
 
 
 class MainConfig(BaseModel):
     data_config: Optional[DataConfig] = Field(
-        default_factory=DataConfig, alias="data", nullable=True
+        default_factory=DataConfig,
+        alias="data",
     )
     """
     Configuration for intermediate data storage and retrieval.
@@ -71,4 +71,4 @@ def load_config(path: Union[str, Path]) -> MainConfig:
     with open(path) as f:
         data = toml.load(f)
         LOG.info(data)
-        return MainConfig.parse_obj(data)
+        return MainConfig.model_validate(data)