fix pyright

bluesky · Feb 17, 2025 · 775527d · 775527d
1 parent 9e45512
commit 775527d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 23 deletions.
diff --git a/src/ophyd_async/tango/core/__init__.py b/src/ophyd_async/tango/core/__init__.py
@@ -12,15 +12,14 @@
 from ._tango_transport import (
     AttributeProxy,
     CommandProxy,
-    DevStateEnum,
     TangoSignalBackend,
     ensure_proper_executor,
     get_dtype_extended,
     get_python_type,
     get_tango_trl,
     get_trl_descriptor,
 )
-from ._utils import get_device_trl_and_attr, get_full_attr_trl
+from ._utils import DevStateEnum, get_device_trl_and_attr, get_full_attr_trl
 
 __all__ = [
     "AttributeProxy",

diff --git a/src/ophyd_async/tango/core/_converters.py b/src/ophyd_async/tango/core/_converters.py
@@ -1,6 +1,7 @@
 from typing import Any, Generic
 
 import numpy as np
+from numpy.typing import NDArray
 
 from ophyd_async.core import (
     SignalDatatypeT,
@@ -33,54 +34,59 @@ def value(self, value: int):
         return self._labels[value]
 
 
-class TangoEnumSpectrumConverter(TangoEnumConverter):
-    def write_value(self, value: np.ndarray[Any, str]):
+class TangoEnumSpectrumConverter(TangoConverter):
+    def __init__(self, labels: list[str]):
+        self._labels = labels
+
+    def write_value(self, value: NDArray[np.str_]) -> NDArray[np.integer]:
         # should return array of ints
         return np.array([self._labels.index(v) for v in value])
 
-    def value(self, value: np.ndarray[Any, int]):
+    def value(self, value) -> NDArray[np.str_]:
         # should return array of strs
         return np.array([self._labels[v] for v in value])
 
 
-class TangoEnumImageConverter(TangoEnumConverter):
-    def write_value(self, value: np.ndarray[Any, str]):
-        # should return array of ints
+class TangoEnumImageConverter(TangoConverter):
+    def __init__(self, labels: list[str]):
+        self._labels = labels
+
+    def write_value(self, value: NDArray[np.str_]) -> NDArray[np.integer]:
         return np.vstack([[self._labels.index(v) for v in row] for row in value])
 
-    def value(self, value: np.ndarray[Any, int]):
-        # should return array of strs
+    def value(self, value) -> NDArray[np.str_]:
         return np.vstack([[self._labels[v] for v in row] for row in value])
 
 
 class TangoDevStateConverter(TangoConverter):
     _labels = [e.value for e in DevStateEnum]
 
-    def write_value(self, value: Any) -> Any:
+    def write_value(self, value: str) -> DevState:
         idx = self._labels.index(value)
         return DevState(idx)
 
-    def value(self, value: DevState) -> Any:
+    def value(self, value: DevState) -> str:
         idx = int(value)
         return self._labels[idx]
 
 
-class TangoDevStateSpectrumConverter(TangoDevStateConverter):
-    def write_value(self, value):
-        # should return array of tango `DevState`s
+class TangoDevStateSpectrumConverter(TangoConverter):
+    _labels = [e.value for e in DevStateEnum]
+
+    def write_value(self, value: NDArray[np.str_]) -> NDArray[DevState]:
         return np.array(
             [DevState(self._labels.index(v)) for v in value], dtype=DevState
         )
 
-    def value(self, value):
-        # should return array of strs
+    def value(self, value: NDArray[DevState]) -> NDArray[np.str_]:
         result = np.array([self._labels[int(v)] for v in value])
         return result
 
 
-class TangoDevStateImageConverter(TangoDevStateConverter):
-    def write_value(self, value):
-        # should return array of tango `DevState`s
+class TangoDevStateImageConverter(TangoConverter):
+    _labels = [e.value for e in DevStateEnum]
+
+    def write_value(self, value: NDArray[np.str_]) -> NDArray[DevState]:
         result = np.vstack(
             [
                 np.array([DevState(self._labels.index(v)) for v in row], dtype=DevState)
@@ -90,6 +96,5 @@ def write_value(self, value):
 
         return result
 
-    def value(self, value):
-        # should return array of strs
+    def value(self, value: NDArray[DevState]) -> NDArray[np.str_]:
         return np.vstack([[self._labels[int(v)] for v in row] for row in value])
diff --git a/src/ophyd_async/tango/core/_tango_transport.py b/src/ophyd_async/tango/core/_tango_transport.py
@@ -749,7 +749,7 @@ async def connect(self, timeout: float) -> None:
             await self._connect_and_store_config(self.read_trl, timeout)
         self.proxies[self.read_trl].set_polling(*self._polling)  # type: ignore
         self.converter = make_converter(self.trl_configs[self.read_trl])
-        self.proxies[self.read_trl].set_converter(self.converter)
+        self.proxies[self.read_trl].set_converter(self.converter)  # type: ignore
         self.descriptor = get_trl_descriptor(
             self.datatype, self.read_trl, self.trl_configs
         )