ci: add more linting steps

.gitea/workflows/ci.yml

@@ -1,8 +1,13 @@
 name: CI
 on:
   push:
+    branches:
+      - main
   pull_request:
-  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 
 jobs:
   lint:
@@ -22,6 +27,23 @@ jobs:
           source .venv/bin/activate
           ruff format --check
 
+      - name: Lint
+        run: |
+          source .venv/bin/activate
+          ruff check
+
+      - name: Checkout pyright diff plugin
+        uses: https://github.com/actions/checkout@v5
+        with:
+          repository: mx/diff_quality_basedpyright
+          path: diff_quality_basedpyright
+
+      - name: Typecheck Diff
+        run: |
+          source .venv/bin/activate
+          uv pip install -e diff_quality_basedpyright
+          diff-quality --violations=basedpyright --fail-under=100
+
   test:
     runs-on: ubuntu-latest
     needs: "lint"

src/aarecommon/config/logger.py

@@ -130,7 +130,7 @@ def find_existing_formatter(
         if h.formatter:
             return h.formatter
 
-    for logger_name, logger in logging.Logger.manager.loggerDict.items():
+    for logger in logging.Logger.manager.loggerDict.values():
         if isinstance(logger, logging.Logger):
             for h in logger.handlers:
                 if h.formatter:
@@ -139,7 +139,7 @@ def find_existing_formatter(
     return logging.Formatter(default_fmt, datefmt=default_datefmt)
 
 
-def attach_to_logger(logger_name: str = "", handler: logging.Handler = None):
+def attach_to_logger(logger_name: str = "", handler: logging.Handler | None = None):
     logging.getLogger(logger_name).addHandler(handler)
 
 

src/aarecommon/config/logger_events.py

@@ -1,7 +1,8 @@
 import functools
 import logging
 import time
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 from aarecommon.models.raster_grid import RasterGridRequest
 from aarecommon.models.rotation_scan import RotationScanRequest

src/aarecommon/errors/exception_handler.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 from typing import ClassVar
 
-from aarecommon.errors.codes import AuthErrorCode
 from aarecommon.config.logger import setup_logger
+from aarecommon.errors.codes import AuthErrorCode
 
 logger = setup_logger("aareDAQ")
 
@@ -216,7 +216,6 @@ class AuthenticationException(AareAuthError):
 
 
 class UserRightsException(AareAuthError):
-    _last_log_ts_by_message: dict[str, float] = {}
     _throttle_window_s = 30.0
 
     def __init__(
@@ -229,6 +228,7 @@ class UserRightsException(AareAuthError):
         critical: bool | None = None,
     ):
         super().__init__(message, critical=critical)
+        self._last_log_ts_by_message: dict[str, float] = {}
         self.message = message
         self.status_code = status_code
         self.headers = headers

src/aarecommon/lc_infer_wrapper.py

@@ -2,26 +2,25 @@ import io
 
 import cv2
 import numpy as np
-from aarecommon.models.beamline import MXBeamline
-from aarecommon.config.beamline import cfg_get, mx_beamline
 from aarelcinfer_client import AuthenticatedClient
 from aarelcinfer_client.api import beam, config, predictions
 from aarelcinfer_client.models import LatestPredictionModel, RuntimeConfigPatchModel
 from PIL import Image
 
+from aarecommon.config.beamline import cfg_get, mx_beamline
+from aarecommon.models.beamline import MXBeamline
+
 
 class AareLCInferWrapper:
     def __init__(self, bl: MXBeamline, secret: str = "1s3ng@rd"):
         if bl == MXBeamline.X10SA or bl == MXBeamline.X06DA:
             host = cfg_get("daq.hardware.aarelc_url")
             if host is None:
-                raise Exception("AareLCInferWrapper: AareLC URL not configured")
-        elif bl == MXBeamline.X06SA:
-            raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}")
-        elif bl == MXBeamline.SIMULATED:
+                raise ValueError("AareLCInferWrapper: AareLC URL not configured")
+        elif bl == MXBeamline.X06SA or bl == MXBeamline.SIMULATED:
             raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}")
         else:
-            raise Exception(f"Unknown beamline {bl}")
+            raise ValueError(f"Unknown beamline {bl}")
 
         self.client = AuthenticatedClient(base_url=host, api_key=secret)
         self.client.headers["X-API-Key"] = secret
@@ -79,7 +78,7 @@ if __name__ == "__main__":
         response = wrapper.update_config(patch)
         print("Config updated!")
         print(f"Model reloaded: {getattr(response, 'model_reloaded', False)}")
-    except Exception as e:
+    except Exception as e:  # noqa
         print(f"Error updating config: {e}")
 
     try:
@@ -134,5 +133,5 @@ if __name__ == "__main__":
         # cv2.waitKey(0)
         # cv2.destroyAllWindows()
 
-    except Exception as e:
+    except Exception as e:  # noqa
         print(f"Error processing prediction bundle: {e}")

src/aarecommon/math/coordinate.py

@@ -1,4 +1,3 @@
-from typing import Optional
 
 import numpy as np
 from pydantic import BaseModel
@@ -75,9 +74,9 @@ class Coordinate(BaseModel):
 
 
 class SmargonCoordinate(BaseModel):
-    sh_mm: Optional[Coordinate] = None  # SH coordinate in Smargon
-    phi_deg: Optional[float] = None
-    chi_deg: Optional[float] = None
+    sh_mm: Coordinate | None = None  # SH coordinate in Smargon
+    phi_deg: float | None = None
+    chi_deg: float | None = None
 
     def eq(self, other: "SmargonCoordinate", tol: float) -> bool:
         return (
@@ -101,8 +100,8 @@ def positive_coords(value: Coordinate) -> Coordinate:
 
 
 class AerotechCoordinate(BaseModel):
-    at_mm: Optional[Coordinate] = None
-    omega_deg: Optional[float] = None
+    at_mm: Coordinate | None = None
+    omega_deg: float | None = None
 
     def eq(self, other: "AerotechCoordinate", tol: float) -> bool:
         return (

src/aarecommon/math/diffraction_geometry.py

@@ -1,5 +1,5 @@
 import math
-from typing import Annotated, Tuple
+from typing import Annotated
 
 from pydantic import BaseModel, Field
 
@@ -8,8 +8,8 @@ class DiffractionGeometry(BaseModel):
     energy_keV: Annotated[float, Field(gt=1.0, lt=100.0)]
     dtz_mm: Annotated[float, Field(gt=10.0, lt=5000.0)]
     pixel_size_mm: Annotated[float, Field(ge=0.05, le=0.5)]
-    beam_center_pxl: Tuple[float, float]
-    detector_size_pxl: Tuple[int, int]
+    beam_center_pxl: tuple[float, float]
+    detector_size_pxl: tuple[int, int]
     detector_description: str
     detector_serial_number: str
     poni_rot1_rad: float
@@ -55,6 +55,7 @@ class DiffractionGeometry(BaseModel):
 if __name__ == "__main__":
     from aare.daq.config import BeamlineConfig
     from aare.devices.jfjoch import JFJochWrapper
+
     from aarecommon.config.beamline import mx_beamline
 
     bl = mx_beamline()

src/aarecommon/math/find_xtal.py

@@ -1,10 +1,11 @@
-from typing import Callable, List, Optional
+from collections.abc import Callable
 
 import numpy as np
+from scipy import ndimage
+
 from aarecommon.config.logger import setup_logger
 from aarecommon.models.models import CrystalSize
 from aarecommon.models.raster_grid import CenterOfMassModel, RasterGridRequest
-from scipy import ndimage
 
 logger = setup_logger("aareDAQ")
 
@@ -32,14 +33,14 @@ def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel |
 
 
 def rebuild_array_from_scan_results(
-    scan_results: List,
+    scan_results: list,
     value_field: str,
-    array_shape: Optional[tuple] = None,
+    array_shape: tuple | None = None,
     nx_field: str = "nx",
     ny_field: str = "ny",
     default_value: float = 0.0,
-    threshold: Optional[float] = None,
-    condition_func: Optional[Callable] = None,
+    threshold: float | None = None,
+    condition_func: Callable | None = None,
     apply_filter_before: bool = True,
 ) -> np.ndarray:
     positions = []
@@ -113,13 +114,13 @@ def rebuild_array_from_scan_results(
 
 
 def create_quality_filtered_array(
-    scan_results: List,
+    scan_results: list,
     value_field: str,
-    min_spots: Optional[int] = None,
-    min_efficiency: Optional[float] = 1.0,
-    min_background: Optional[float] = None,
-    exclude_ice: Optional[bool] = True,
-    min_low_res_spots: Optional[float] = 10.0,
+    min_spots: int | None = None,
+    min_efficiency: float | None = 1.0,
+    min_background: float | None = None,
+    exclude_ice: bool | None = True,
+    min_low_res_spots: float | None = 10.0,
     **kwargs,
 ) -> np.ndarray:
     """
@@ -154,10 +155,7 @@ def create_quality_filtered_array(
         if result.spots_low_res < min_background:
             return False
 
-        if result.efficiency < min_efficiency:
-            return False
-
-        return True
+        return not result.efficiency < min_efficiency
 
     # Filter results first
     filtered_results = []
@@ -247,7 +245,7 @@ def get_xtal_size(crystal_size, result_array, r: RasterGridRequest):
     return crystal_size
 
 
-def get_best_b_factor(result_list: List):
+def get_best_b_factor(result_list: list):
     if not result_list:
         return None
     best_b_factor = min(
@@ -259,7 +257,7 @@ def get_best_b_factor(result_list: List):
     return best_b_factor.b
 
 
-def get_best_res(result_list: List):
+def get_best_res(result_list: list):
     if not result_list:
         return None
     best_res = min(
@@ -272,10 +270,7 @@ def get_best_res(result_list: List):
 
 
 def com_nan_check(com):
-    if np.isnan(com.n_x) or np.isnan(com.n_y):
-        return False
-    else:
-        return True
+    return not (np.isnan(com.n_x) or np.isnan(com.n_y))
 
 
 def get_result_list_from_com(images, com: CenterOfMassModel):
@@ -365,7 +360,7 @@ def has_sufficient_low_res_spots(result_array: np.ndarray, min_spots_low_res: fl
 
 
 def compute_crystal_score_array(
-    scan_results: List, w_bkg: float = 0.25, w_low_res: float = 0.75, w_indexed: float = 0.00
+    scan_results: list, w_bkg: float = 0.25, w_low_res: float = 0.75, w_indexed: float = 0.00
 ) -> np.ndarray:
     """Combine bkg (25%), spots_low_res (75%), and spots_indexed (00%) into a 0–100 score.
 
@@ -397,7 +392,7 @@ def _draw_panel(
     mask: np.ndarray,
     label: str,
     threshold: float,
-    grid_size_mm: Optional[tuple[float, float]] = None,
+    grid_size_mm: tuple[float, float] | None = None,
     cbar_label: str = "spots_low_res",
 ) -> None:
     """Shared helper: heatmap + crystal contour + max-cell square on one Axes.
@@ -407,6 +402,7 @@ def _draw_panel(
     """
     import matplotlib.patches as mpatches
     import matplotlib.pyplot as plt
+
     from aarecommon.coordinate import Coordinate
 
     n_nx, n_ny = arr.shape
@@ -574,9 +570,9 @@ CRYSTAL_METHOD_MAP: dict = {
 
 
 def compare_crystal_methods(
-    results: List,
-    arr: Optional[np.ndarray] = None,
-    grid_size_mm: Optional[tuple[float, float]] = None,
+    results: list,
+    arr: np.ndarray | None = None,
+    grid_size_mm: tuple[float, float] | None = None,
 ) -> None:
     """Plot each crystal-detection method side-by-side for visual comparison.
 
@@ -613,10 +609,10 @@ def compare_crystal_methods(
 
 
 def compare_crystal_methods_scored(
-    results: List,
-    score_arr: Optional[np.ndarray] = None,
-    method: Optional[str] = None,
-    grid_size_mm: Optional[tuple[float, float]] = None,
+    results: list,
+    score_arr: np.ndarray | None = None,
+    method: str | None = None,
+    grid_size_mm: tuple[float, float] | None = None,
     w_bkg: float = 0.20,
     w_low_res: float = 0.60,
     w_indexed: float = 0.20,

src/aarecommon/math/sample_geometry.py

@@ -1,9 +1,10 @@
 from typing import Annotated
 
 import numpy as np
-from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
 from pydantic import AfterValidator, BaseModel, Field
 
+from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
+
 
 class SampleGeometryModel(BaseModel):
     # Beam location in the camera coordinates

src/aarecommon/math/simulate_raster.py

@@ -27,13 +27,13 @@ Seed  Scenario
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import List
 
 import numpy as np
-from aarecommon.models.raster_grid import RasterGridRequest
 from jfjoch_client import ScanResult
 from jfjoch_client.models.scan_result_images_inner import ScanResultImagesInner
 
+from aarecommon.models.raster_grid import RasterGridRequest
+
 # ---------------------------------------------------------------------------
 # Cluster geometry descriptor
 # ---------------------------------------------------------------------------
@@ -59,7 +59,7 @@ class _ClusterParams:
 
 @dataclass
 class _SeedConfig:
-    clusters: List[_ClusterParams] = field(default_factory=list)
+    clusters: list[_ClusterParams] = field(default_factory=list)
 
 
 # ---------------------------------------------------------------------------

src/aarecommon/models/aerotech.py

@@ -1,8 +1,8 @@
 from enum import Enum
-from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
 
 from aarecommon.models.rotation_scan import RotationScanRequest
-from pydantic import BaseModel, ConfigDict, Field
 
 
 class TaskEnum(Enum):
@@ -53,10 +53,10 @@ class AerotechAxisStatus(BaseModel):
 
 class AerotechStatus(BaseModel):
     state: str
-    x: Optional[AerotechAxisStatus] = None
-    y: Optional[AerotechAxisStatus] = None
-    z: Optional[AerotechAxisStatus] = None
-    u: Optional[AerotechAxisStatus] = None
+    x: AerotechAxisStatus | None = None
+    y: AerotechAxisStatus | None = None
+    z: AerotechAxisStatus | None = None
+    u: AerotechAxisStatus | None = None
 
     model_config = ConfigDict(extra="allow")
 
@@ -133,10 +133,10 @@ class AerotechStatus(BaseModel):
 
 
 class AerotechTarget(BaseModel):
-    x: Optional[float] = None
-    y: Optional[float] = None
-    z: Optional[float] = None
-    u: Optional[float] = None
+    x: float | None = None
+    y: float | None = None
+    z: float | None = None
+    u: float | None = None
 
     def to_payload(self) -> dict:
         return self.model_dump(exclude_none=True)

src/aarecommon/models/auth.py

@@ -1,4 +1,5 @@
 from enum import Enum
+
 from pydantic import BaseModel
 
 
@@ -57,8 +58,8 @@ class BatonStatus(BaseModel):
 
 
 def get_user():
-    import os
     import getpass
+    import os
 
     try:
         return os.getlogin()

src/aarecommon/models/automation.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from enum import Enum
 from typing import Any, Literal
 
@@ -56,7 +56,7 @@ class AutomationProgress:
     ) -> None:
         self.events.append(
             LogEvent(
-                ts=datetime.now(timezone.utc),
+                ts=datetime.now(UTC),
                 level=level,
                 code=code,
                 exception_class=exception_class,

src/aarecommon/models/models.py

@@ -1,15 +1,16 @@
 import re
 from dataclasses import dataclass
 from enum import Enum
-from typing import Annotated, List, Literal, Optional, Tuple
+from typing import Annotated, Literal
+
+from jfjoch_client.models.scan_result import ScanResult
+from pydantic import AfterValidator, AliasChoices, BaseModel, ConfigDict, Field, field_validator
 
-from aarecommon.models.beamline import MXBeamline
 from aarecommon.math.coordinate import Coordinate, positive_coords
 from aarecommon.math.diffraction_geometry import DiffractionGeometry
 from aarecommon.math.sample_geometry import SampleGeometryModel
+from aarecommon.models.beamline import MXBeamline
 from aarecommon.models.tell import TellStateModel
-from jfjoch_client.models.scan_result import ScanResult
-from pydantic import AfterValidator, AliasChoices, BaseModel, ConfigDict, Field, field_validator
 
 
 class StagePositionEnum(Enum):
@@ -21,7 +22,7 @@ class StagePositionEnum(Enum):
 
 class TokenData(BaseModel):
     sub: str  # Username
-    pgroups: List[str]
+    pgroups: list[str]
     session: int
     staff: bool = False
 
@@ -42,7 +43,7 @@ class PuckInfo(BaseModel):
     puck_name: str
     dewar_name: str
     user: str = ""
-    location: Optional[DewarAddress] = None
+    location: DewarAddress | None = None
 
 
 # From TELL to database after loading
@@ -54,43 +55,41 @@ class PuckLoadedInfo(BaseModel):
 class DataCollectionParameters(BaseModel):
     model_config = ConfigDict(from_attributes=True, populate_by_name=True)
 
-    directory: Optional[str] = None
-    oscillation: Optional[float] = None  # Only accept positive float
-    exposure: Optional[float] = None  # Only accept positive floats between 0 and 1
-    totalangle: Optional[int] = Field(  # was totalrange
+    directory: str | None = None
+    oscillation: float | None = None  # Only accept positive float
+    exposure: float | None = None  # Only accept positive floats between 0 and 1
+    totalangle: int | None = Field(  # was totalrange
         default=None, validation_alias=AliasChoices("totalangle", "totalrange")
     )  # Only accept positive integers between 0 and 360
-    transmission: Optional[int] = None  # Only accept positive integers between 0 and 100
-    targetresolution: Optional[float] = None  # Only accept positive float
-    beamsize: Optional[str] = None
-    aperture: Optional[int] = None  # Optional string field
-    datacollectiontype: Optional[str] = (
+    transmission: int | None = None  # Only accept positive integers between 0 and 100
+    targetresolution: float | None = None  # Only accept positive float
+    beamsize: str | None = None
+    aperture: int | None = None  # Optional string field
+    datacollectiontype: str | None = (
         None  # Only accept "standard", other types might be added later
     )
-    processingpipeline: Optional[str] = ""  # Only accept "gopy", "autoproc", "xia2dials"
-    spacegroupnumber: Optional[int] = None  # Only accept positive integers between 1 and 230
-    unitcell: Optional[str] = Field(  # was cellparameters
+    processingpipeline: str | None = ""  # Only accept "gopy", "autoproc", "xia2dials"
+    spacegroupnumber: int | None = None  # Only accept positive integers between 1 and 230
+    unitcell: str | None = Field(  # was cellparameters
         default=None, validation_alias=AliasChoices("unitcell", "cellparameters")
     )  # Must be a set of six positive floats or integers
-    rescutkey: Optional[str] = None  # Only accept "is" or "cchalf"
-    rescutvalue: Optional[float] = None  # Must be a positive float if rescutkey is provided
-    processingresolution: Optional[float] = Field(  # was userresolution
+    rescutkey: str | None = None  # Only accept "is" or "cchalf"
+    rescutvalue: float | None = None  # Must be a positive float if rescutkey is provided
+    processingresolution: float | None = Field(  # was userresolution
         default=None, validation_alias=AliasChoices("processingresolution", "userresolution")
     )
-    pdbid: Optional[str] = (
-        ""  # Accepts either the format of the protein data bank code or {provided}
-    )
-    autoprocfull: Optional[bool] = None
-    procfull: Optional[bool] = None
-    adpenabled: Optional[bool] = None
-    noano: Optional[bool] = None
-    ffcscampaign: Optional[bool] = None
-    trustedhigh: Optional[float] = None  # Should be a float between 0 and 2.0
-    autoprocextraparams: Optional[str] = None  # Optional string field
-    chiphiangles: Optional[float] = None  # Optional float field between 0 and 30
-    dose: Optional[float] = None  # Optional float field
+    pdbid: str | None = ""  # Accepts either the format of the protein data bank code or {provided}
+    autoprocfull: bool | None = None
+    procfull: bool | None = None
+    adpenabled: bool | None = None
+    noano: bool | None = None
+    ffcscampaign: bool | None = None
+    trustedhigh: float | None = None  # Should be a float between 0 and 2.0
+    autoprocextraparams: str | None = None  # Optional string field
+    chiphiangles: float | None = None  # Optional float field between 0 and 30
+    dose: float | None = None  # Optional float field
     cloud: bool = True
-    pdbmodel: Optional[str] = None
+    pdbmodel: str | None = None
 
     @field_validator("directory", mode="after")
     @classmethod
@@ -175,12 +174,12 @@ class SampleShortInfo(BaseModel):
     dewar_name: str
     sample_name: str
     run_number: int
-    aaredb_params: Optional[DataCollectionParameters] = None
+    aaredb_params: DataCollectionParameters | None = None
     user: str = ""
     pin: Annotated[int, Field(ge=1, le=16)]
     location: DewarAddress | None = None
-    priority: Optional[float] = 1.0
-    comment: Optional[str] = None
+    priority: float | None = 1.0
+    comment: str | None = None
     mount_count: int = 0
     rotation_count: int = 0
     raster_count: int = 0
@@ -207,7 +206,7 @@ class SampleShortInfo(BaseModel):
 
 
 class SampleShortInfoList(BaseModel):
-    s: List[SampleShortInfo]
+    s: list[SampleShortInfo]
 
 
 class BeamMarkCoeffModel(BaseModel):
@@ -217,8 +216,8 @@ class BeamMarkCoeffModel(BaseModel):
     Default is beam center in 1000,1000 at any zoom level
     """
 
-    coeff_x: Tuple[float, float, float] = (0, 0, 1000.0)
-    coeff_y: Tuple[float, float, float] = (0, 0, 1000.0)
+    coeff_x: tuple[float, float, float] = (0, 0, 1000.0)
+    coeff_y: tuple[float, float, float] = (0, 0, 1000.0)
 
     def apply(self, zoom: float):
         return Coordinate(
@@ -263,11 +262,11 @@ class MLBoxModel(BaseModel):
 
     @staticmethod
     def from_tuple(
-        cls: MLBoxType, box_tuple: tuple[float, float, float, float], conf: float
+        klass: MLBoxType, box_tuple: tuple[float, float, float, float], conf: float
     ) -> "MLBoxModel":
         x1, y1, x2, y2 = box_tuple
         return MLBoxModel(
-            cls=cls,
+            cls=klass,
             box=BoundingBoxModel(
                 top_x=float(x1), top_y=float(y1), bottom_x=float(x2), bottom_y=float(y2)
             ),
@@ -279,18 +278,18 @@ class MLOutputModel(BaseModel):
     boxes: dict[str, MLBoxModel] = {}
 
     @staticmethod
-    def get_class_str(cls: MLBoxType) -> str:
-        if cls == MLBoxType.LOOP_ALL:
+    def get_class_str(klass: MLBoxType) -> str:
+        if klass == MLBoxType.LOOP_ALL:
             return "Loop_all"
-        if cls == MLBoxType.PIN:
+        if klass == MLBoxType.PIN:
             return "Pin"
-        if cls == MLBoxType.CRYSTAL:
+        if klass == MLBoxType.CRYSTAL:
             return "Crystal"
-        if cls == MLBoxType.LOOP_FACE:
+        if klass == MLBoxType.LOOP_FACE:
             return "Loop_face"
-        if cls == MLBoxType.ICE:
+        if klass == MLBoxType.ICE:
             return "Ice"
-        if cls == MLBoxType.NEEDLE:
+        if klass == MLBoxType.NEEDLE:
             return "Needle"
         return "Unknown"
 
@@ -483,11 +482,11 @@ def def_zoom(beamline) -> ZoomModel:
                 1000: SampleCameraSettings(gain=0, exposure=0.25),
             }
         )
-    elif beamline == MXBeamline.X10SA:
-        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
-    elif beamline == MXBeamline.X06SA:
-        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
-    elif beamline == MXBeamline.SIMULATED:
+    elif (
+        beamline == MXBeamline.X10SA
+        or beamline == MXBeamline.X06SA
+        or beamline == MXBeamline.SIMULATED
+    ):
         return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
     else:
         raise ValueError(f"Invalid beamline: {beamline}")
@@ -507,11 +506,11 @@ def def_bl_zoom(beamline) -> ZoomModel:
                 1000: SampleCameraSettings(gain=0, exposure=0.25),
             }
         )
-    elif beamline == MXBeamline.X10SA:
-        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
-    elif beamline == MXBeamline.X06SA:
-        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
-    elif beamline == MXBeamline.SIMULATED:
+    elif (
+        beamline == MXBeamline.X10SA
+        or beamline == MXBeamline.X06SA
+        or beamline == MXBeamline.SIMULATED
+    ):
         return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
     else:
         raise ValueError(f"Invalid beamline: {beamline}")
@@ -527,9 +526,7 @@ def def_loop_centering_zoom(beamline) -> ZoomModel:
                 280: SampleCameraSettings(gain=0, exposure=0.002),
             }
         )
-    elif beamline == MXBeamline.X06SA:
-        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
-    elif beamline == MXBeamline.SIMULATED:
+    elif beamline == MXBeamline.X06SA or beamline == MXBeamline.SIMULATED:
         return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
     else:
         raise ValueError(f"Invalid beamline: {beamline}")
@@ -663,11 +660,11 @@ class SimpleStrategyInputModel(BaseModel):
     incr_omega_deg: float = 0.2
     d_vis: float = 1.5
     current_temp_k: float = 100
-    filename: Optional[str] = None
+    filename: str | None = None
 
 
 class SimpleScanParameters(BaseModel):
-    filename: Optional[str] = None
+    filename: str | None = None
     dtz: float = 150
     exp_time_s: float = 0.02
     start_omega_deg: float = 0

src/aarecommon/models/raster_grid.py

@@ -1,11 +1,12 @@
-from typing import Annotated, List
+from typing import Annotated
 
 import numpy as np
-from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
-from aarecommon.math.sample_geometry import SampleGeometryModel
 from jfjoch_client.models.scan_result import ScanResult
 from pydantic import AfterValidator, BaseModel, Field
 
+from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords
+from aarecommon.math.sample_geometry import SampleGeometryModel
+
 
 class RasterGridRequest(BaseModel):
     dtz: float | None = None
@@ -76,7 +77,7 @@ class CompletedRasterGridElem(BaseModel):
 
 
 class CompletedRasterGrid(BaseModel):
-    r: List[CompletedRasterGridElem]
+    r: list[CompletedRasterGridElem]
 
 
 class RasterPayloadModel(BaseModel):
@@ -85,7 +86,7 @@ class RasterPayloadModel(BaseModel):
     sample_id: int
     attach_image: bool = True
     centre_of_mass: CenterOfMassModel | None = None
-    raster_score: List[float | None] | None = None
+    raster_score: list[float | None] | None = None
     center_pxl: Coordinate | None
     start_pxl: Coordinate
     cell_size_pxl: Annotated[Coordinate, AfterValidator(positive_coords)]

src/aarecommon/models/rotation_scan.py

@@ -1,7 +1,8 @@
-from aarecommon.math.coordinate import SmargonCoordinate
 from jfjoch_client.models.scan_result import ScanResult
 from pydantic import BaseModel
 
+from aarecommon.math.coordinate import SmargonCoordinate
+
 
 class RotationScanRequest(BaseModel):
     dtz: float | None = None

src/aarecommon/recurrence_watcher.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, TypeAlias
+from typing import TypeAlias
 
 from aarecommon.errors import exception_handler
 from aarecommon.errors.exception_handler import (

tests/test_autofocus_tools.py

@@ -1,4 +1,5 @@
 import numpy as np
+
 from aarecommon.math.autofocus import focus_measure_blob_size, focus_measure_edges
 
 

tests/test_beamline.py

@@ -8,10 +8,9 @@ from aarecommon.models.beamline import MXBeamline
 
 
 def test_mx_beamline_default():
-    with patch.dict(os.environ, {}, clear=True):
-        with pytest.raises(ValueError) as e:
-            # If BEAMLINE is not set, it should raise
-            mx_beamline()
+    with patch.dict(os.environ, {}, clear=True), pytest.raises(ValueError) as e:
+        # If BEAMLINE is not set, it should raise
+        mx_beamline()
     assert e.match("set the BEAMLINE")
 
 

tests/test_coordinate.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+
 from aarecommon.math.coordinate import (
     AerotechCoordinate,
     Coordinate,

tests/test_data_collection_parameters.py

@@ -1,7 +1,8 @@
 import pytest
-from aarecommon.models.models import DataCollectionParameters
 from pydantic import ValidationError
 
+from aarecommon.models.models import DataCollectionParameters
+
 
 def test_directory_defaults_when_missing():
     params = DataCollectionParameters()

tests/test_diffraction_geometry.py

@@ -1,4 +1,5 @@
 import pytest
+
 from aarecommon.math.diffraction_geometry import DiffractionGeometry
 
 

tests/test_exception_handler.py

@@ -1,4 +1,5 @@
 import pytest
+
 from aarecommon.errors.exception_handler import (
     AareDBCommunicationError,
     AerotechCommunicationError,

tests/test_find_xtal.py

@@ -2,6 +2,7 @@ from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
+
 from aarecommon.math.find_xtal import (
     compute_crystal_score_array,
     create_quality_filtered_array,

tests/test_logger_events.py

@@ -3,6 +3,7 @@ import time
 from unittest.mock import MagicMock
 
 import pytest
+
 from aarecommon.config.logger_events import log_timing, merge_log_context
 
 

tests/test_raster_grid_common.py

@@ -1,4 +1,5 @@
 import pytest
+
 from aarecommon.math.raster_grid import grid_to_image_id, image_id_to_grid
 
 

tests/test_zoom_model_camera_settings.py

@@ -1,4 +1,5 @@
 import pytest
+
 from aarecommon.models.models import SampleCameraSettings, ZoomModel