Run pydantic migrator on murfey and fix issues that arise

Author: stephen-riggs

pyproject.toml

@@ -31,7 +31,8 @@ classifiers = [
 ]
 dependencies = [
     "defusedxml", # For safely parsing XML files
-    "pydantic<2", # Locked to <2 by zocalo
+    "pydantic>=2",
+    "pydantic_settings",
     "requests",
     "rich",
     "werkzeug",

src/murfey/client/instance_environment.py

@@ -8,7 +8,8 @@
 from typing import Callable, Dict, List, NamedTuple, Optional, Set
 from urllib.parse import ParseResult
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic_core.core_schema import FieldValidationInfo
 
 from murfey.client.watchdir import DirWatcher
 
@@ -40,20 +41,20 @@ class MurfeyInstanceEnvironment(BaseModel):
     destination_registry: Dict[str, str] = {}
     watchers: Dict[Path, DirWatcher] = {}
     demo: bool = False
-    data_collection_group_ids: Dict[str, int] = {}
-    data_collection_ids: Dict[str, int] = {}
-    processing_job_ids: Dict[str, Dict[str, int]] = {}
-    autoproc_program_ids: Dict[str, Dict[str, int]] = {}
     id_tag_registry: Dict[str, List[str]] = {
         "data_collection_group": [],
         "data_collection": [],
         "processing_job": [],
         "auto_proc_program": [],
     }
+    listeners: Dict[str, Set[Callable]] = {}
+    data_collection_group_ids: Dict[str, int] = {}
+    data_collection_ids: Dict[str, int] = {}
+    processing_job_ids: Dict[str, Dict[str, int]] = {}
+    autoproc_program_ids: Dict[str, Dict[str, int]] = {}
     data_collection_parameters: dict = {}
     movies: Dict[Path, MovieTracker] = {}
     motion_corrected_movies: Dict[Path, List[str]] = {}
-    listeners: Dict[str, Set[Callable]] = {}
     movie_tilt_pair: Dict[Path, str] = {}
     tilt_angles: Dict[str, List[List[str]]] = {}
     movie_counters: Dict[str, itertools.count] = {}
@@ -64,42 +65,41 @@ class MurfeyInstanceEnvironment(BaseModel):
     murfey_session: Optional[int] = None
     samples: Dict[Path, SampleInfo] = {}
 
-    class Config:
-        validate_assignment: bool = True
-        arbitrary_types_allowed: bool = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    @validator("data_collection_group_ids")
-    def dcg_callback(cls, v, values):
+    @field_validator("data_collection_group_ids")
+    def dcg_callback(cls, v, info: FieldValidationInfo):
         with global_env_lock:
-            for l in values.get("listeners", {}).get("data_collection_group_ids", []):
+            for l in info.data.get("listeners", {}).get(
+                "data_collection_group_ids", []
+            ):
                 for k in v.keys():
-                    if k not in values["id_tag_registry"]["data_collection"]:
+                    if k not in info.data["id_tag_registry"]["data_collection"]:
                         l(k)
         return v
 
-    @validator("data_collection_ids")
-    def dc_callback(cls, v, values):
+    @field_validator("data_collection_ids")
+    def dc_callback(cls, v, info: FieldValidationInfo):
         with global_env_lock:
-            for l in values.get("listeners", {}).get("data_collection_ids", []):
+            for l in info.data.get("listeners", {}).get("data_collection_ids", []):
                 for k in v.keys():
-                    if k not in values["id_tag_registry"]["processing_job"]:
+                    if k not in info.data["id_tag_registry"]["processing_job"]:
                         l(k)
         return v
 
-    @validator("processing_job_ids")
-    def job_callback(cls, v, values):
+    @field_validator("processing_job_ids")
+    def job_callback(cls, v, info: FieldValidationInfo):
         with global_env_lock:
-            for l in values.get("listeners", {}).get("processing_job_ids", []):
+            for l in info.data.get("listeners", {}).get("processing_job_ids", []):
                 for k in v.keys():
-                    if k not in values["id_tag_registry"]["auto_proc_program"]:
+                    if k not in info.data["id_tag_registry"]["auto_proc_program"]:
                         l(k, v[k]["ispyb-relion"])
         return v
 
-    @validator("autoproc_program_ids")
-    def app_callback(cls, v, values):
-        # logger.info(f"autoproc program ids validator: {v}")
+    @field_validator("autoproc_program_ids")
+    def app_callback(cls, v, info: FieldValidationInfo):
         with global_env_lock:
-            for l in values.get("listeners", {}).get("autoproc_program_ids", []):
+            for l in info.data.get("listeners", {}).get("autoproc_program_ids", []):
                 for k in v.keys():
                     if v[k].get("em-tomo-preprocess"):
                         l(k, v[k]["em-tomo-preprocess"])

src/murfey/client/tui/screens.py

@@ -197,7 +197,7 @@ def validate_form(form: dict, model: BaseModel) -> bool:
     try:
         convert = lambda x: None if x == "None" else x
         validated = model(**{k: convert(v) for k, v in form.items()})
-        log.info(validated.dict())
+        log.info(validated.model_dump())
         return True
     except (AttributeError, ValidationError) as e:
         log.warning(f"Form validation failed: {str(e)}")

src/murfey/server/config.py

@@ -7,20 +7,21 @@
 from typing import Dict, List, Optional, Union
 
 import yaml
-from pydantic import BaseModel, BaseSettings
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
 
 
 class MachineConfig(BaseModel):
     acquisition_software: List[str]
-    calibrations: Dict[str, Dict[str, Union[dict, float]]]
+    calibrations: Dict[str, Dict[int, Union[dict, float]]]
     data_directories: Dict[Path, str]
     rsync_basepath: Path
     murfey_db_credentials: str
     crypto_key: str
     default_model: Path
     display_name: str = ""
     image_path: Optional[Path] = None
-    software_versions: Dict[str, str] = {}
+    software_versions: Dict[str, float] = {}
     external_executables: Dict[str, str] = {}
     external_executables_eer: Dict[str, str] = {}
     external_environment: Dict[str, str] = {}

src/murfey/server/demo_api.py

@@ -14,7 +14,8 @@
 from fastapi.responses import FileResponse, HTMLResponse
 from ispyb.sqlalchemy import BLSession
 from PIL import Image
-from pydantic import BaseModel, BaseSettings
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
 from sqlalchemy import func
 from sqlmodel import col, select
 from werkzeug.utils import secure_filename

src/murfey/server/main.py

@@ -8,7 +8,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from prometheus_client import make_asgi_app
-from pydantic import BaseSettings
+from pydantic_settings import BaseSettings
 
 import murfey.server
 import murfey.server.api.auth

src/murfey/util/models.py

@@ -219,13 +219,13 @@ class SPAProcessFile(BaseModel):
     tag: str
     path: str
     description: str
-    processing_job: Optional[int]
-    data_collection_id: Optional[int]
+    processing_job: Optional[int] = None
+    data_collection_id: Optional[int] = None
     image_number: int
-    autoproc_program_id: Optional[int]
-    foil_hole_id: Optional[int]
-    pixel_size: Optional[float]
-    dose_per_frame: Optional[float]
+    autoproc_program_id: Optional[int] = None
+    foil_hole_id: Optional[int] = None
+    pixel_size: Optional[float] = None
+    dose_per_frame: Optional[float] = None
     mc_binning: Optional[int] = 1
     gain_ref: Optional[str] = None
     extract_downscale: bool = True
@@ -236,7 +236,7 @@ class SPAProcessFile(BaseModel):
 class ProcessingParametersSPA(BaseModel):
     tag: str
     dose_per_frame: float
-    gain_ref: Optional[str]
+    gain_ref: Optional[str] = None
     experiment_type: str
     voltage: float
     image_size_x: int
@@ -247,12 +247,12 @@ class ProcessingParametersSPA(BaseModel):
     acquisition_software: str
     use_cryolo: bool
     symmetry: str
-    mask_diameter: Optional[int]
-    boxsize: Optional[int]
+    mask_diameter: Optional[int] = None
+    boxsize: Optional[int] = None
     downscale: bool
-    small_boxsize: Optional[int]
+    small_boxsize: Optional[int] = None
     eer_fractionation: int
-    particle_diameter: Optional[float]
+    particle_diameter: Optional[float] = None
     magnification: Optional[int] = None
     total_exposed_dose: Optional[float] = None
     c2aperture: Optional[float] = None
@@ -261,14 +261,14 @@ class ProcessingParametersSPA(BaseModel):
     phase_plate: bool = False
 
     class Base(BaseModel):
-        dose_per_frame: Optional[float]
-        gain_ref: Optional[str]
+        dose_per_frame: Optional[float] = None
+        gain_ref: Optional[str] = None
         use_cryolo: bool
         symmetry: str
-        mask_diameter: Optional[int]
-        boxsize: Optional[int]
+        mask_diameter: Optional[int] = None
+        boxsize: Optional[int] = None
         downscale: bool
-        small_boxsize: Optional[int]
+        small_boxsize: Optional[int] = None
         eer_fractionation: int