Migrate to pydantic version 2

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 ]
 dependencies = [
     "defusedxml", # For safely parsing XML files
-    "pydantic<2", # Locked to <2 by zocalo
+    "pydantic>=2",
     "requests",
     "rich",
     "werkzeug",
@@ -41,13 +41,12 @@ cicd = [
     "pytest-cov", # Used by Azure Pipelines for PyTest coverage reports
 ]
 client = [
-    "procrunner",
     "textual==0.42.0",
     "websocket-client",
     "xmltodict",
 ]
 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
@@ -58,7 +57,7 @@ server = [
     "backports.entry_points_selectable",
     "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",
@@ -71,7 +70,7 @@ server = [
     "sqlmodel",
     "stomp-py<=8.1.0", # 8.1.1 (released 2024-04-06) doesn't work with our project
     "uvicorn[standard]",
-    "zocalo",
+    "zocalo>=1",
 ]
 [project.urls]
 Bug-Tracker = "https://github.com/DiamondLightSource/python-murfey/issues"

src/murfey/cli/transfer.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 import argparse
+import subprocess
 from pathlib import Path
 from urllib.parse import urlparse
 
-import procrunner
 import requests
 from rich.console import Console
 from rich.prompt import Confirm
@@ -78,6 +78,6 @@ def run():
         cmd.extend(list(Path(args.source or ".").glob("*")))
         cmd.append(f"{murfey_url.hostname}::{args.destination}")
 
-    result = procrunner.run(cmd)
+    result = subprocess.run(cmd)
     if result.returncode:
         console.print(f"[red]rsync failed returning code {result.returncode}")

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 ValidationInfo
 
 from murfey.client.watchdir import DirWatcher
 
@@ -41,20 +42,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] = {}
@@ -65,42 +66,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: ValidationInfo):
         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: ValidationInfo):
         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: ValidationInfo):
         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: ValidationInfo):
         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/rsync.py

@@ -18,8 +18,6 @@
 from typing import Callable, List, NamedTuple
 from urllib.parse import ParseResult
 
-import procrunner
-
 from murfey.client.tui.status_bar import StatusBar
 from murfey.util import Observer
 
@@ -433,28 +431,29 @@ def parse_stderr(line: str):
         result: subprocess.CompletedProcess | None = None
         success = True
         if rsync_stdin:
-            result = procrunner.run(
+            result = subprocess.run(
                 rsync_cmd,
-                callback_stdout=parse_stdout,
-                callback_stderr=parse_stderr,
-                working_directory=str(self._basepath),
-                stdin=rsync_stdin,
-                print_stdout=False,
-                print_stderr=False,
+                cwd=str(self._basepath),
+                capture_output=True,
+                input=rsync_stdin,
             )
-            success = result.returncode == 0 if result else False
+            for stdout_line in result.stdout.decode("utf8", "replace").split("\n"):
+                parse_stdout(stdout_line)
+            for stderr_line in result.stderr.decode("utf8", "replace").split("\n"):
+                parse_stderr(stderr_line)
+            success = result.returncode == 0
 
         if rsync_stdin_remove:
             rsync_cmd.insert(-2, "--remove-source-files")
-            result = procrunner.run(
+            result = subprocess.run(
                 rsync_cmd,
-                callback_stdout=parse_stdout,
-                callback_stderr=parse_stderr,
-                working_directory=str(self._basepath),
-                stdin=rsync_stdin_remove,
-                print_stdout=False,
-                print_stderr=False,
+                cwd=str(self._basepath),
+                input=rsync_stdin_remove,
             )
+            for stdout_line in result.stdout.decode("utf8", "replace").split("\n"):
+                parse_stdout(stdout_line)
+            for stderr_line in result.stderr.decode("utf8", "replace").split("\n"):
+                parse_stderr(stderr_line)
 
             if success:
                 success = result.returncode == 0 if result else False

src/murfey/client/tui/app.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import logging
+import subprocess
 from datetime import datetime
 from functools import partial
 from pathlib import Path
 from queue import Queue
 from typing import Awaitable, Callable, Dict, List, OrderedDict, TypeVar
 from urllib.parse import urlparse
 
-import procrunner
 import requests
 from textual.app import App
 from textual.reactive import reactive
@@ -203,7 +203,7 @@ def _start_rsyncer(
         if self._environment:
             self._environment.default_destinations[source] = destination
             if self._environment.gain_ref and visit_path:
-                gain_rsync = procrunner.run(
+                gain_rsync = subprocess.run(
                     [
                         "rsync",
                         str(self._environment.gain_ref),

src/murfey/client/tui/screens.py

@@ -2,6 +2,7 @@
 
 # import contextlib
 import logging
+import subprocess
 from datetime import datetime
 from functools import partial
 from pathlib import Path
@@ -17,7 +18,6 @@
     TypeVar,
 )
 
-import procrunner
 import requests
 from pydantic import BaseModel, ValidationError
 from rich.box import SQUARE
@@ -199,7 +199,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)}")
@@ -846,7 +846,7 @@ def on_button_pressed(self, event):
             if self.app._environment.demo:
                 log.info(f"Would perform {' '.join(cmd)}")
             else:
-                gain_rsync = procrunner.run(cmd)
+                gain_rsync = subprocess.run(cmd)
                 if gain_rsync.returncode:
                     log.warning(
                         f"Gain reference file {self._dir_tree._gain_reference} was not successfully transferred to {visit_path}/processing"

src/murfey/instrument_server/api.py

@@ -143,7 +143,7 @@ def start_multigrid_watcher(
         demo=True,
         do_transfer=True,
         processing_enabled=not watcher_spec.skip_existing_processing,
-        _machine_config=watcher_spec.configuration.dict(),
+        _machine_config=watcher_spec.configuration.model_dump(),
         token=tokens.get(session_id, "token"),
         data_collection_parameters=data_collection_parameters.get(label, {}),
     )
@@ -156,7 +156,7 @@ def start_multigrid_watcher(
         (watcher_spec.source / d).mkdir(exist_ok=True)
     watchers[session_id] = MultigridDirWatcher(
         watcher_spec.source,
-        watcher_spec.configuration.dict(),
+        watcher_spec.configuration.model_dump(),
         skip_existing_processing=watcher_spec.skip_existing_processing,
     )
     watchers[session_id].subscribe(controllers[session_id]._start_rsyncer_multigrid)
@@ -221,7 +221,7 @@ 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():
         data_collection_parameters[proc_param_block.label][k] = v
     return {"success": True}
 

src/murfey/server/api/spa.py

@@ -20,15 +20,17 @@ def _cryolo_model_path(visit: str, instrument_name: str) -> Path:
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
-    if machine_config.model_search_directory:
+    if machine_config.picking_model_search_directory:
         visit_directory = (
             machine_config.rsync_basepath
             / (machine_config.rsync_module or "data")
             / str(datetime.now().year)
             / visit
         )
         possible_models = list(
-            (visit_directory / machine_config.model_search_directory).glob("*.h5")
+            (visit_directory / machine_config.picking_model_search_directory).glob(
+                "*.h5"
+            )
         )
         if possible_models:
             return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1]

src/murfey/server/demo_api.py

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

@@ -7,7 +7,8 @@
 from typing import Dict, List, Literal, Optional, Union
 
 import yaml
-from pydantic import BaseModel, BaseSettings
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
 
 
 class MachineConfig(BaseModel):
@@ -57,7 +58,7 @@ class MachineConfig(BaseModel):
     upstream_data_download_directory: Optional[Path] = None  # Set by microscope config
     upstream_data_tiff_locations: List[str] = ["processed"]  # Location of CLEM TIFFs
 
-    model_search_directory: str = "processing"
+    picking_model_search_directory: str = "processing"
     initial_model_search_directory: str = "processing/initial_model"
 
     failure_queue: str = ""

src/murfey/util/db.py

@@ -757,7 +757,8 @@ def setup(url: str):
 
 def clear(url: str):
     engine = create_engine(url)
-    metadata = sqlalchemy.MetaData(engine)
-    metadata.reflect()
+    metadata = sqlalchemy.MetaData()
+    metadata.create_all(engine)
+    metadata.reflect(engine)
 
     metadata.drop_all(engine)