From 6483d50d1d0ae0dc788cd3c541ff97935c31d8b5 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 8 Sep 2026 15:44:47 +0200 Subject: [PATCH] style: ignore pyright issues in redis return values --- src/aare/daq/auth.py | 4 +- src/aare/daq/config.py | 107 +++++++++++++++++++---------------------- src/aare/daq/daq.py | 2 +- src/aare/daq/server.py | 9 ++-- 4 files changed, 58 insertions(+), 64 deletions(-) diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index aeadb149..68fd2ddc 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -99,8 +99,8 @@ def authenticate_user(cfg: BeamlineConfig, username: str) -> str: def parse_token(token: str = Depends(oauth2_scheme)) -> TokenData: try: payload = jwt.decode(token, jwt_key(), algorithms=[ALGORITHM]) - token = TokenData(**payload) - return token + _token = TokenData(**payload) + return _token except jwt.PyJWTError as e: raise AuthenticationException( message="Invalid token", diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 60902e67..14db0716 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -4,9 +4,9 @@ import json import time from dataclasses import asdict, is_dataclass from datetime import datetime +from typing import Any import numpy as np -import redis from aarecommon.config.beamline import cfg_get from aarecommon.config.logger import setup_logger from aarecommon.math.coordinate import AerotechCoordinate, Coordinate @@ -37,6 +37,7 @@ from aarecommon.models.models import ( ZoomModel, zoom_manager, ) +from redis.client import PubSub, Redis from redis.lock import Lock as RedisLock from aare.daq.config_model import LocalContactConfigModel @@ -99,7 +100,8 @@ class BeamlineConfig: host = "localhost" else: host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch") - self.redis = redis.Redis(host=host, port=6379, db=0, decode_responses=True) + self.hw_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") + self.redis = Redis(host=host, port=6379, db=0, decode_responses=True) self.simulated_detector = bl is MXBeamline.SIMULATED self._initialize_optional_yaml_defaults() @@ -144,7 +146,7 @@ class BeamlineConfig: def _current_gui_session_ttl(self, session: int) -> int | None: try: - ttl = int(self.redis.ttl(self._gui_session_key(session))) + ttl = int(self.redis.ttl(self._gui_session_key(session))) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except Exception: logger.debug("Could not read the GUI session TTL", exc_info=True) return None @@ -158,7 +160,7 @@ class BeamlineConfig: if raw is None: return None try: - return OpenGuiSessionInfo(**json.loads(raw)) + return OpenGuiSessionInfo(**json.loads(raw)) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except Exception: logger.warning(f"Failed to parse GUI session info for session {session}", exc_info=True) return None @@ -239,7 +241,7 @@ class BeamlineConfig: self.purge_expired_gui_sessions() def purge_expired_gui_sessions(self) -> None: - session_ids = self.redis.smembers(self._gui_sessions_index_key()) + session_ids: set[str] = self.redis.smembers(self._gui_sessions_index_key()) # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 if not session_ids: return @@ -297,7 +299,7 @@ class BeamlineConfig: self.redis.delete(f"{self._bl}:allow_non_staff_request_from_staff") def generate_session(self) -> int: - return int(self.redis.incr(f"{self._bl}:session")) + return int(self.redis.incr(f"{self._bl}:session")) # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 @property def active_session(self) -> int | None: @@ -368,7 +370,7 @@ class BeamlineConfig: # Redis compatibility: # some Redis servers do not support EXPIRE with GT option. # Emulate "extend only if greater" manually while holding the lock. - if current_ttl is None or current_ttl < 0 or current_ttl < expiry_sec: + if current_ttl is None or float(current_ttl) < 0 or float(current_ttl) < expiry_sec: # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 self.redis.expire(key, expiry_sec) else: raise RuntimeError( @@ -458,10 +460,7 @@ class BeamlineConfig: def can_transfer_baton_now(self) -> bool: """Check if baton can be transferred (beamline not mid-operation).""" - # Can't transfer while beamline is busy - # Add automation queue check here when you implement it - # return not (self.state_busy or self.automation_queue_running) - raise NotImplementedError("Checking for busy state is not implemented") + return not self.hw_lock.locked() def execute_baton_transfer( self, @@ -508,8 +507,7 @@ class BeamlineConfig: @property def pgroup(self) -> str | None: - tmp = self.redis.get(f"{self._bl}:pgroup") - return tmp + return self.redis.get(f"{self._bl}:pgroup") # pyright: ignore[reportReturnType] # using sync client - remove ignore on upgrade to redis v8 @pgroup.setter def pgroup(self, pgroup: str | None) -> None: @@ -520,8 +518,7 @@ class BeamlineConfig: @property def commissioning_mode(self) -> bool: - tmp = self.redis.get(f"{self._bl}:commissioning_mode") - return tmp is not None + return self.redis.get(f"{self._bl}:commissioning_mode") is not None @commissioning_mode.setter def commissioning_mode(self, commisioning_mode: bool) -> None: @@ -545,8 +542,8 @@ class BeamlineConfig: if raw_value is None: return BeamlineStateEnum.Maintenance try: - int_value = int(raw_value) # Ensure it's an integer - return BeamlineStateEnum(int_value) # Convert to BeamlineStateEnum + int_value = int(raw_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return BeamlineStateEnum(int_value) except (ValueError, KeyError): raise ValueError( f"Invalid 'mx_state' value: {raw_value}. Expected integer corresponding to a BeamlineStateEnum." @@ -558,12 +555,12 @@ class BeamlineConfig: @property def tell_mount_count(self) -> int: - return int(self.redis.incr(f"{self._bl}:tell_mount_count")) + return int(self.redis.incr(f"{self._bl}:tell_mount_count")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def pixel_to_mm(self, zoom: float) -> float: cfg = self.settings base_pixel_in_mm = 1.0 / ( - cfg.camera_translation_factor_b * np.exp(cfg.camera_translation_factor_a * zoom) + cfg.camera_translation_factor_b * np.exp(cfg.camera_translation_factor_a * zoom) # pyright: ignore[reportOptionalOperand] ) # Apply lens magnification correction relative to the default 10x lens. # A lower magnification lens (e.g. 5x) makes each pixel cover more physical space. @@ -600,17 +597,9 @@ class BeamlineConfig: @property def beam_size_mm(self) -> Coordinate: - tmp_x = self.redis.get(f"{self._bl}:beam_size_x") - tmp_y = self.redis.get(f"{self._bl}:beam_size_y") - if tmp_x: - val_x = float(tmp_x) - else: - val_x = 0.04 - if tmp_y: - val_y = float(tmp_y) - else: - val_y = 0.04 - return Coordinate(x=val_x, y=val_y) + x = float(self.redis.get(f"{self._bl}:beam_size_x") or 0.4) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + y = float(self.redis.get(f"{self._bl}:beam_size_y") or 0.4) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return Coordinate(x=x, y=y) @beam_size_mm.setter def beam_size_mm(self, data: Coordinate): @@ -622,7 +611,7 @@ class BeamlineConfig: if tmp is None: return BeamlineSettingsModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 model = BeamlineSettingsModel(**data_dict) return model @@ -644,7 +633,7 @@ class BeamlineConfig: if tmp is None: return CryojetSettingsModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return CryojetSettingsModel(**data_dict) @cryojet_settings.setter @@ -652,7 +641,7 @@ class BeamlineConfig: self.redis.set(f"{self._bl}:cryojet_settings", data.model_dump_json()) def get_alc_bkg(self, zoom: float, exp: float, gain: float) -> np.ndarray | None: - return base64_to_numpy(self.redis.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) + return base64_to_numpy(self.redis.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def put_alc_bkg(self, zoom: float, exp: float, gain: float, data: np.ndarray): self.redis.set(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}", numpy_to_base64(data)) @@ -663,7 +652,7 @@ class BeamlineConfig: if tmp is None: return SampleShortInfoList(s=[]) - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfoList(**data_dict) def spreadsheet_pgroup(self, pgroups: list[str]) -> SampleShortInfoList: @@ -675,7 +664,7 @@ class BeamlineConfig: def spreadsheet(self, data: SampleShortInfoList): self.redis.set(f"{self._bl}:sample_spreadsheet", data.model_dump_json()) - def listen_changes_spreadsheet(self) -> redis.client.PubSub: + def listen_changes_spreadsheet(self) -> PubSub: self.redis.config_set("notify-keyspace-events", "KEA") pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:sample_spreadsheet") @@ -687,7 +676,7 @@ class BeamlineConfig: if tmp is None: return SampleShortInfoList(s=[]) - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfoList(**data_dict) @reference_tools.setter @@ -706,7 +695,7 @@ class BeamlineConfig: if tmp is None: return None - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfo(**data_dict) @current_sample.setter @@ -737,8 +726,8 @@ class BeamlineConfig: print("no zoom mode given, defaulting to user mode") return ZoomModeEnum.User try: - int_value = int(raw_value) # Ensure it's an integer - return ZoomModeEnum(int_value) # Convert to ZoomModeEnum + int_value = int(raw_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return ZoomModeEnum(int_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except (ValueError, KeyError): raise ValueError( f"Invalid 'mx_zoom_mode' value: {raw_value}. Expected integer corresponding to a ZoomModeEnum." @@ -788,9 +777,9 @@ class BeamlineConfig: mode = mode or self.zoom_mode key = f"{self._bl}:{self.zoom_setting_string(mode)}" tmp = self.redis.get(key) - model = ZoomModel(**json.loads(tmp)) if tmp is not None else zoom_manager(mode, self._mxb) - model.z[zoom_value] = settings - self.redis.set(key, model.model_dump_json()) + model = ZoomModel(**json.loads(tmp)) if tmp is not None else zoom_manager(mode, self._mxb) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + model.z[zoom_value] = settings # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + self.redis.set(key, model.model_dump_json()) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @property def abr_meas_pos(self) -> AerotechCoordinate: @@ -798,7 +787,7 @@ class BeamlineConfig: if tmp is None: return ABR_POS_MOUNT - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return AerotechCoordinate(**data_dict) @abr_meas_pos.setter @@ -810,7 +799,7 @@ class BeamlineConfig: tmp = self.redis.get(f"{self._bl}:dtz") if tmp is None: return None - return float(tmp) + return float(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @dtz.setter def dtz(self, dtz: float): @@ -821,7 +810,7 @@ class BeamlineConfig: tmp = self.redis.get(f"{self._bl}:dtz_safe_position") if tmp is None: return None - return float(tmp) + return float(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @dtz_safe_position.setter def dtz_safe_position(self, dtz: float): @@ -851,11 +840,11 @@ class BeamlineConfig: ) vals = self.redis.hgetall(f"{self._bl}:beam_mark") - if len(vals) >= 3: + if len(vals) >= 3: # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 zooms = [] x_pxls = [] y_pxls = [] - for k, v in vals.items(): + for k, v in vals.items(): # pyright: ignore[reportArgumentType,reportAttributeAccessIssue] # using sync client - remove ignore on upgrade to redis v8 zooms.append(float(k)) x_pxls.append(float(json.loads(v)["x"])) y_pxls.append(float(json.loads(v)["y"])) @@ -917,10 +906,10 @@ class BeamlineConfig: def get_mount_failure_streak(self) -> int: value = self.redis.get(self._mount_failure_streak_key()) - return int(value) if value else 0 + return int(value) if value else 0 # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def increment_mount_failure_streak(self) -> int: - return int(self.redis.incr(self._mount_failure_streak_key())) + return int(self.redis.incr(self._mount_failure_streak_key())) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def reset_mount_failure_streak(self) -> None: self.redis.delete(self._mount_failure_streak_key()) @@ -979,16 +968,18 @@ class BeamlineConfig: self.redis.set(self._automation_progress_seq_key(), 0) self.redis.delete(self._automation_progress_key()) - def get_automation_progress_state(self) -> dict: + def get_automation_progress_state(self) -> dict[str, Any]: seq_raw = self.redis.get(self._automation_progress_seq_key()) payload_raw = self.redis.get(self._automation_progress_key()) - seq = int(seq_raw) if seq_raw is not None else 0 - progress = json.loads(payload_raw) if payload_raw else None + seq = int(seq_raw) if seq_raw is not None else 0 # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + progress = json.loads(payload_raw) if payload_raw else None # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return {"seq": seq, "progress": progress} - def set_automation_progress_state(self, progress: AutomationProgress | dict) -> dict: + def set_automation_progress_state( + self, progress: AutomationProgress | dict[str, Any] + ) -> dict[str, Any]: def _json_default(value): if isinstance(value, datetime): return value.isoformat() @@ -1003,7 +994,7 @@ class BeamlineConfig: else: raise TypeError(f"Unsupported automation progress type: {type(progress).__name__}") - next_seq = int(self.redis.incr(self._automation_progress_seq_key())) + next_seq = int(self.redis.incr(self._automation_progress_seq_key())) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 self.redis.set( self._automation_progress_key(), json.dumps(payload, separators=(",", ":"), default=_json_default), @@ -1017,7 +1008,7 @@ class BeamlineConfig: if tmp is None: return 0 try: - return int(tmp) + return int(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except (TypeError, ValueError): logger.warning( "Failed Mount Count is not an integer, resetting to 0.", @@ -1033,7 +1024,7 @@ class BeamlineConfig: self.redis.set(f"{self._bl}:failed_mount_count", count) def increment_failed_mount_count(self) -> int: - return int(self.redis.incr(f"{self._bl}:failed_mount_count")) + return int(self.redis.incr(f"{self._bl}:failed_mount_count")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def _runtime_sim_key(self, name: str) -> str: return f"{self._bl}:runtime:simulate:{name}" @@ -1132,7 +1123,7 @@ class BeamlineConfig: logger.warning(f"Failed to read detector metadata from Redis: {e}", exc_info=True) return {} - def set_detector_metadata(self, payload: dict) -> dict: + def set_detector_metadata(self, payload: dict[str, Any]) -> dict[str, Any]: safe_payload = dict(payload or {}) safe_payload["dtz_low"] = self._coerce_optional_float(safe_payload.get("dtz_low")) safe_payload["dtz_high"] = self._coerce_optional_float(safe_payload.get("dtz_high")) @@ -1144,7 +1135,7 @@ class BeamlineConfig: return safe_payload @property - def cached_detector_metadata(self) -> dict: + def cached_detector_metadata(self) -> dict[str, Any]: return self.get_detector_metadata() @property diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 47570a32..4f301e4e 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -2190,7 +2190,7 @@ class AareDAQ: sample=self.sample, sample_geometry=self.sample_geometry, filename=filename, - upload_image=self._aare.upload_image, + upload_image=self._aare.upload_image, # pyright: ignore[reportArgumentType] logger=logger, max_images=self.AUTO_RASTER_MAX_IMAGES, min_cell_size_mm=self.AUTO_RASTER_MAX_IMAGES, diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 7f6c192f..b7c39029 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import functools import hmac @@ -7,6 +9,7 @@ import os import time from collections.abc import AsyncGenerator from contextlib import asynccontextmanager +from typing import Any import uvicorn from aarecommon.config.beamline import mx_beamline @@ -123,7 +126,7 @@ async def lifespan(application: FastAPI): # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── bl = mx_beamline() cfg = BeamlineConfig(bl) - hardware_busy_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") + hardware_busy_lock = cfg.hw_lock daq = AareDAQ(cfg, bl, bl_dispatch, hardware_busy_lock) cfg.state = daq.read_current_state_from_bec() @@ -655,7 +658,7 @@ async def anneal(time_s: float, token: str = Depends(oauth2_scheme)): @app.post("/smargon/initialize") @needs_hw_lock -async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: +async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Initialise Smargon. Staff only. @@ -673,7 +676,7 @@ async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: @app.post("/bec/load_user_macros") @needs_hw_lock -async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict: +async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Load BEC user macros. Staff only. """