Migrate to pydantic version 2

pyproject.toml

@@ -32,7 +32,8 @@ classifiers = [
 dependencies = [
     "backports.entry_points_selectable",
     "defusedxml", # For safely parsing XML files
-    "pydantic<2", # Locked to <2 by cygwin terminal
+    "pydantic>=2",
+    "pydantic-settings",
     "requests",
     "rich",
     "werkzeug",
@@ -47,7 +48,7 @@ client = [
     "websocket-client",
 ]
 developer = [
-    "bump-my-version<0.11.0", # Version control
+    "bump-my-version", # Version control
     "ipykernel", # Enable interactive coding with VS Code and Jupyter Notebook
     "pre-commit", # Formatting, linting, type checking, etc.
     "pytest", # Test code functionality
@@ -61,7 +62,7 @@ server = [
     "aiohttp",
     "cryptography",
     "fastapi[standard]",
-    "ispyb", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python; v10.0.0: sqlalchemy <2, mysql-connector-python >=8.0.32
+    "ispyb>=10.2.4", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python;
     "jinja2",
     "mrcfile",
     "numpy<2",
@@ -73,7 +74,7 @@ server = [
     "sqlalchemy[postgresql]", # Add as explicit dependency
     "sqlmodel",
     "stomp-py<=8.1.0", # 8.1.1 (released 2024-04-06) doesn't work with our project
-    "zocalo",
+    "zocalo>=1",
 ]
 [project.urls]
 Bug-Tracker = "https://github.com/DiamondLightSource/python-murfey/issues"

src/murfey/client/contexts/spa.py

@@ -506,7 +506,8 @@ def post_transfer(
                         except Exception as e:
                             # try to continue if position information gathering fails so that movie is processed anyway
                             logger.warning(
-                                f"Unable to register foil hole for {str(file_transferred_to)}. Exception: {str(e)}"
+                                f"Unable to register foil hole for {str(file_transferred_to)}. Exception: {str(e)}",
+                                exc_info=True,
                             )
                             foil_hole = None
 

src/murfey/client/contexts/tomo.py

@@ -173,7 +173,7 @@
                         )
 
         except Exception as e:
-            logger.error(f"ERROR {e}, {self.data_collection_parameters}")
+            logger.error(f"ERROR {e}, {self.data_collection_parameters}", exc_info=True)
 
     def _file_transferred_to(
         self, environment: MurfeyInstanceEnvironment, source: Path, file_path: Path
@@ -533,7 +533,9 @@
                     try:
                         for_parsing = xml.read()
                     except Exception:
-                        logger.warning(f"Failed to parse file {metadata_file}")
+                        logger.warning(
+                            f"Failed to parse file {metadata_file}", exc_info=True
+                        )
                         return OrderedDict({})
                     data = xmltodict.parse(for_parsing)
                 try:
@@ -628,7 +630,9 @@
                     / int(mdoc_metadata["num_eer_frames"])
                 )
         except Exception as e:
-            logger.error(f"Exception encountered in metadata gathering: {str(e)}")
+            logger.error(
+                f"Exception encountered in metadata gathering: {str(e)}", exc_info=True
+            )
             return OrderedDict({})
 
         return mdoc_metadata

src/murfey/client/instance_environment.py

@@ -8,7 +8,7 @@
 from typing import Dict, List, NamedTuple, Optional
 from urllib.parse import ParseResult
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from murfey.client.watchdir import DirWatcher
 
@@ -56,9 +56,7 @@ class MurfeyInstanceEnvironment(BaseModel):
     samples: Dict[Path, SampleInfo] = {}
     rsync_url: str = ""
 
-    class Config:
-        validate_assignment: bool = True
-        arbitrary_types_allowed: bool = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def clear(self):
         self.sources = []

src/murfey/client/tui/screens.py

@@ -196,7 +196,7 @@
     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/instrument_server/api.py

@@ -335,14 +335,14 @@ def register_processing_parameters(
     session_id: MurfeySessionID, proc_param_block: ProcessingParameterBlock
 ):
     data_collection_parameters[proc_param_block.label] = {}
-    for k, v in proc_param_block.params.dict().items():
+    for k, v in proc_param_block.params.model_dump().items():
         if v is not None:
             data_collection_parameters[proc_param_block.label][k] = v
     if controllers.get(session_id):
         controllers[session_id].data_collection_parameters.update(
             data_collection_parameters[proc_param_block.label]
         )
-        for k, v in proc_param_block.params.dict().items():
+        for k, v in proc_param_block.params.model_dump().items():
             if v is not None and hasattr(controllers[session_id]._environment, k):
                 setattr(controllers[session_id]._environment, k, v)
     return {"success": True}

src/murfey/server/api/clem.py

@@ -10,7 +10,7 @@
 
 from backports.entry_points_selectable import entry_points
 from fastapi import APIRouter
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 from sqlalchemy.exc import NoResultFound
 from sqlmodel import Session, select
 
@@ -820,10 +820,8 @@ class AlignAndMergeParams(BaseModel):
     flatten: Literal["mean", "min", "max", ""] = ""
     align_across: Literal["enabled", ""] = ""
 
-    @validator(
-        "images",
-        pre=True,
-    )
+    @field_validator("images", mode="before")
+    @classmethod
     def parse_stringified_list(cls, value):
         if isinstance(value, str):
             try:

src/murfey/server/api/prometheus.py

@@ -39,7 +39,7 @@ def increment_rsync_file_count(
         logger.error(
             f"Failed to find rsync instance for visit {sanitise(visit_name)} "
             "with the following properties: \n"
-            f"{rsyncer_info.dict()}",
+            f"{rsyncer_info.model_dump()}",
             exc_info=True,
         )
         return None

src/murfey/server/api/session_control.py

@@ -96,7 +96,7 @@ def get_current_visits(instrument_name: str, db=ispyb_db):
 
 
 class SessionInfo(BaseModel):
-    session_id: Optional[int]
+    session_id: Optional[int] = None
     session_name: str = ""
     rescale: bool = True
 

src/murfey/server/api/workflow.py

@@ -81,7 +81,7 @@ class DCGroupParameters(BaseModel):
     tag: str
     atlas: str = ""
     sample: Optional[int] = None
-    atlas_pixel_size: int = 0
+    atlas_pixel_size: float = 0
 
 
 @router.post("/visits/{visit_name}/{session_id}/register_data_collection_group")
@@ -346,13 +346,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
@@ -608,9 +608,9 @@ class TomoProcessFile(BaseModel):
     tag: str
     image_number: int
     pixel_size: float
-    dose_per_frame: Optional[float]
+    dose_per_frame: Optional[float] = None
     frame_count: int
-    tilt_axis: Optional[float]
+    tilt_axis: Optional[float] = None
     mc_uuid: Optional[int] = None
     voltage: float = 300
     mc_binning: int = 1
@@ -894,7 +894,7 @@ class Sample(BaseModel):
     sample_group_id: int
     sample_id: int
     subsample_id: int
-    image_path: Optional[Path]
+    image_path: Optional[Path] = None
 
 
 @correlative_router.get("/visit/{visit_name}/samples")

src/murfey/server/demo_api.py

@@ -13,32 +13,38 @@
 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 select
 from werkzeug.utils import secure_filename
 
 import murfey.server.api.bootstrap
 import murfey.server.prometheus as prom
-from murfey.server import (
-    _flush_grid_square_records,
-    _murfey_id,
-    get_hostname,
-    get_microscope,
-    sanitise,
-    sanitise_path,
-)
-from murfey.server import shutdown as _shutdown
-from murfey.server import templates
-from murfey.server.api.auth import MurfeySessionID, validate_token
+from murfey.server.api import templates
+from murfey.server.api.auth import MurfeySessionIDFrontend as MurfeySessionID
+from murfey.server.api.auth import validate_token
 from murfey.server.api.session_info import Visit
 from murfey.server.api.workflow import (
     DCGroupParameters,
     DCParameters,
     ProcessingJobParameters,
 )
+from murfey.server.feedback import (
+    _flush_grid_square_records,
+    _murfey_id,
+    get_microscope,
+    sanitise,
+)
 from murfey.server.murfey_db import murfey_db
-from murfey.util.config import MachineConfig, from_file, security_from_file
+from murfey.server.run import shutdown as _shutdown
+from murfey.util import sanitise_path
+from murfey.util.config import (
+    MachineConfig,
+    from_file,
+    get_hostname,
+    security_from_file,
+)
 from murfey.util.db import (
     AutoProcProgram,
     ClientEnvironment,

src/murfey/server/main.py

@@ -7,7 +7,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/config.py

@@ -4,11 +4,21 @@
 import socket
 from functools import lru_cache
 from pathlib import Path
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional
 
 import yaml
 from backports.entry_points_selectable import entry_points
-from pydantic import BaseModel, BaseSettings, Extra, validator
+from pydantic import BaseModel, ConfigDict, RootModel, ValidationInfo, field_validator
+from pydantic_settings import BaseSettings
+
+
+class MagnificationTable(RootModel[dict[int, float]]):
+    pass
+
+
+CALIBRATIONS_VALIDATION_SCHEMAS = {
+    "magnification": MagnificationTable,
+}
 
 
 class MachineConfig(BaseModel):  # type: ignore
@@ -26,7 +36,7 @@
     # Hardware and software -----------------------------------------------------------
     camera: str = "FALCON"
     superres: bool = False
-    calibrations: dict[str, dict[str, Union[dict, float]]]
+    calibrations: dict[str, Any]
     acquisition_software: list[str]
     software_versions: dict[str, str] = {}
     software_settings_output_directories: dict[str, list[str]] = {}
@@ -72,7 +82,7 @@
 
     # Particle picking setup
     default_model: Path
-    model_search_directory: str = "processing"
+    picking_model_search_directory: str = "processing"
     initial_model_search_directory: str = "processing/initial_model"
 
     # Data analysis plugins
@@ -93,15 +103,43 @@
     node_creator_queue: str = "node_creator"
     notifications_queue: str = "pato_notification"
 
-    class Config:
-        """
-        Inner class that defines this model's parsing and serialising behaviour
-        """
+    # Pydantic BaseModel settings
+    model_config = ConfigDict(extra="allow")
+
+    @field_validator("calibrations", mode="before")
+    @classmethod
+    def validate_calibration_data(
+        cls, v: dict[str, dict[Any, Any]]
+    ) -> dict[str, dict[Any, Any]]:
+        # Pass the calibration dictionaries through their matching Pydantic models, if any are set
+        if isinstance(v, dict):
+            validated = {}
+            for (
+                key,
+                value,
+            ) in v.items():
+                model_cls = CALIBRATIONS_VALIDATION_SCHEMAS.get(key)
+                if model_cls:
+                    try:
+                        # Validate and store as a dict object with the corrected types
+                        validated[key] = model_cls.model_validate(value).root
+                    except Exception as e:
+                        raise ValueError(f"Validation failed for key '{key}': {e}")
+                else:
+                    validated[key] = value
+            return validated
+        # Let it validate and fail as-is
+        return v
 
-        extra = Extra.allow
-        json_encoders = {
-            Path: str,
-        }
+    @field_validator("software_versions", mode="before")
+    @classmethod
+    def validate_software_versions(cls, v: dict[str, Any]) -> dict[str, str]:
+        # Software versions should be numerical strings, even if they appear int- or float-like
+        if isinstance(v, dict):
+            validated = {key: str(value) for key, value in v.items()}
+            return validated
+        # Let it validate and fail as-is
+        return v
 
 
 def from_file(config_file_path: Path, instrument: str = "") -> dict[str, MachineConfig]:
@@ -144,16 +182,13 @@
     graylog_host: str = ""
     graylog_port: Optional[int] = None
 
-    class Config:
-        json_encoders = {
-            Path: str,
-        }
+    model_config = ConfigDict()
 
-    @validator("graylog_port")
+    @field_validator("graylog_port")
     def check_port_present_if_host_is(
-        cls, v: Optional[int], values: dict, **kwargs
+        cls, v: Optional[int], info: ValidationInfo, **kwargs
     ) -> Optional[int]:
-        if values["graylog_host"] and v is None:
+        if info.data.get("graylog_host") and v is None:
             raise ValueError("The Graylog port must be set if the Graylog host is")
         return v