Fix pyright

Author: GDYendell

pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
     "aioserial",
     "numpy",
     "pydantic",
-    "pvi~=0.9.0",
+    "pvi~=0.10.0",
     "pytango",
     "softioc",
 ]

src/fastcs/backends/epics/gui.py

@@ -7,11 +7,11 @@
     LED,
     ButtonPanel,
     ComboBox,
-    Component,
+    ComponentUnion,
     Device,
     Grid,
     Group,
-    ReadWidget,
+    ReadWidgetUnion,
     SignalR,
     SignalRW,
     SignalW,
@@ -22,7 +22,7 @@
     TextWrite,
     ToggleButton,
     Tree,
-    WriteWidget,
+    WriteWidgetUnion,
 )
 from pydantic import ValidationError
 
@@ -56,7 +56,7 @@ def _get_pv(self, attr_path: list[str], name: str):
         return f"{attr_prefix}:{name.title().replace('_', '')}"
 
     @staticmethod
-    def _get_read_widget(attribute: AttrR) -> ReadWidget:
+    def _get_read_widget(attribute: AttrR) -> ReadWidgetUnion:
         match attribute.datatype:
             case Bool():
                 return LED()
@@ -68,7 +68,7 @@ def _get_read_widget(attribute: AttrR) -> ReadWidget:
                 raise FastCSException(f"Unsupported type {type(datatype)}: {datatype}")
 
     @staticmethod
-    def _get_write_widget(attribute: AttrW) -> WriteWidget:
+    def _get_write_widget(attribute: AttrW) -> WriteWidgetUnion:
         match attribute.allowed_values:
             case allowed_values if allowed_values is not None:
                 return ComboBox(choices=allowed_values)
@@ -87,7 +87,7 @@ def _get_write_widget(attribute: AttrW) -> WriteWidget:
 
     def _get_attribute_component(
         self, attr_path: list[str], name: str, attribute: Attribute
-    ):
+    ) -> SignalR | SignalW | SignalRW:
         pv = self._get_pv(attr_path, name)
         name = name.title().replace("_", "")
 
@@ -108,6 +108,8 @@ def _get_attribute_component(
             case AttrW():
                 write_widget = self._get_write_widget(attribute)
                 return SignalW(name=name, write_pv=pv, write_widget=write_widget)
+            case _:
+                raise FastCSException(f"Unsupported attribute type: {type(attribute)}")
 
     def _get_command_component(self, attr_path: list[str], name: str):
         pv = self._get_pv(attr_path, name)
@@ -136,8 +138,8 @@ def create_gui(self, options: EpicsGUIOptions | None = None) -> None:
         formatter = DLSFormatter()
         formatter.format(device, options.output_path)
 
-    def extract_mapping_components(self, mapping: SingleMapping) -> list[Component]:
-        components: Tree[Component] = []
+    def extract_mapping_components(self, mapping: SingleMapping) -> Tree:
+        components: Tree = []
         attr_path = mapping.controller.path
 
         for name, sub_controller in mapping.controller.get_sub_controllers().items():
@@ -151,7 +153,7 @@ def extract_mapping_components(self, mapping: SingleMapping) -> list[Component]:
                 )
             )
 
-        groups: dict[str, list[Component]] = {}
+        groups: dict[str, list[ComponentUnion]] = {}
         for attr_name, attribute in mapping.attributes.items():
             try:
                 signal = self._get_attribute_component(

src/fastcs/backends/epics/ioc.py

@@ -162,7 +162,7 @@ async def async_record_set(value: T):
             record.set(enum_value_to_index(attribute, value))
     else:
 
-        async def async_record_set(value: T):  # type: ignore
+        async def async_record_set(value: T):
             record.set(value)
 
     record = _get_input_record(f"{pv_prefix}:{pv_name}", attribute)
@@ -173,8 +173,10 @@ async def async_record_set(value: T):  # type: ignore
 
 def _get_input_record(pv: str, attribute: AttrR) -> RecordWrapper:
     if attr_is_enum(attribute):
-        # https://github.com/python/mypy/issues/16789
-        state_keys = dict(zip(MBB_STATE_FIELDS, attribute.allowed_values, strict=False))  # type: ignore
+        assert attribute.allowed_values is not None and all(
+            isinstance(v, str) for v in attribute.allowed_values
+        )
+        state_keys = dict(zip(MBB_STATE_FIELDS, attribute.allowed_values, strict=False))
         return builder.mbbIn(pv, **state_keys)
 
     match attribute.datatype:
@@ -210,7 +212,7 @@ async def async_write_display(value: T):
         async def on_update(value):
             await attribute.process_without_display_update(value)
 
-        async def async_write_display(value: T):  # type: ignore
+        async def async_write_display(value: T):
             record.set(value, process=False)
 
     record = _get_output_record(
@@ -223,7 +225,10 @@ async def async_write_display(value: T):  # type: ignore
 
 def _get_output_record(pv: str, attribute: AttrW, on_update: Callable) -> Any:
     if attr_is_enum(attribute):
-        state_keys = dict(zip(MBB_STATE_FIELDS, attribute.allowed_values, strict=False))  # type: ignore
+        assert attribute.allowed_values is not None and all(
+            isinstance(v, str) for v in attribute.allowed_values
+        )
+        state_keys = dict(zip(MBB_STATE_FIELDS, attribute.allowed_values, strict=False))
         return builder.mbbOut(pv, always_update=True, on_update=on_update, **state_keys)
 
     match attribute.datatype:

src/fastcs/backends/tango/dsr.py

@@ -144,7 +144,7 @@ def _collect_dev_properties(mapping: Mapping) -> dict[str, Any]:
 
 def _collect_dev_init(mapping: Mapping) -> dict[str, Callable]:
     async def init_device(tango_device: Device):
-        await server.Device.init_device(tango_device)
+        await server.Device.init_device(tango_device)  # type: ignore
         tango_device.set_state(DevState.ON)
         await mapping.controller.connect()
 

src/fastcs/connections/ip_connection.py

@@ -12,43 +12,59 @@ class IPConnectionSettings:
     port: int = 25565
 
 
-class IPConnection:
-    def __init__(self):
-        self._reader, self._writer = (None, None)
+@dataclass
+class StreamConnection:
+    reader: asyncio.StreamReader
+    writer: asyncio.StreamWriter
+
+    def __post_init__(self):
         self._lock = asyncio.Lock()
 
-    async def connect(self, settings: IPConnectionSettings):
-        self._reader, self._writer = await asyncio.open_connection(
-            settings.ip, settings.port
-        )
+    async def __aenter__(self):
+        await self._lock.acquire()
+        return self
 
-    def ensure_connected(self):
-        if self._reader is None or self._writer is None:
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._lock.release()
+
+    async def send_message(self, message) -> None:
+        self.writer.write(message.encode("utf-8"))
+        await self.writer.drain()
+
+    async def receive_response(self) -> str:
+        data = await self.reader.readline()
+        return data.decode("utf-8")
+
+    async def close(self):
+        self.writer.close()
+        await self.writer.wait_closed()
+
+
+class IPConnection:
+    def __init__(self):
+        self.__connection = None
+
+    @property
+    def _connection(self) -> StreamConnection:
+        if self.__connection is None:
             raise DisconnectedError("Need to call connect() before using IPConnection.")
 
+        return self.__connection
+
+    async def connect(self, settings: IPConnectionSettings):
+        reader, writer = await asyncio.open_connection(settings.ip, settings.port)
+        self.__connection = StreamConnection(reader, writer)
+
     async def send_command(self, message) -> None:
-        async with self._lock:
-            self.ensure_connected()
-            await self._send_message(message)
+        async with self._connection as connection:
+            await connection.send_message(message)
 
     async def send_query(self, message) -> str:
-        async with self._lock:
-            self.ensure_connected()
-            await self._send_message(message)
-            return await self._receive_response()
+        async with self._connection as connection:
+            await connection.send_message(message)
+            return await connection.receive_response()
 
-    # TODO: Figure out type hinting for connections. TypeGuard fails to work as expected
     async def close(self):
-        async with self._lock:
-            self.ensure_connected()
-            self._writer.close()
-            await self._writer.wait_closed()
-            self._reader, self._writer = (None, None)
-
-    async def _send_message(self, message) -> None:
-        self._writer.write(message.encode("utf-8"))
-        await self._writer.drain()
-
-    async def _receive_response(self) -> str:
-        data = await self._reader.readline()
-        return data.decode("utf-8")
+        async with self._connection as connection:
+            await connection.close()
+            self.__connection = None

src/fastcs/connections/serial_connection.py

@@ -20,33 +20,33 @@ def __init__(self):
         self._lock = asyncio.Lock()
 
     async def connect(self, settings: SerialConnectionSettings) -> None:
-        self.stream = aioserial.AioSerial(port=settings.port, baudrate=settings.baud)
+        self.__stream = aioserial.AioSerial(port=settings.port, baudrate=settings.baud)
 
-    def ensure_open(self):
-        if self.stream is None:
+    @property
+    def _stream(self) -> aioserial.AioSerial:
+        if self.__stream is None:
             raise NotOpenedError(
                 "Need to call connect() before using SerialConnection."
             )
 
+        return self.__stream
+
     async def send_command(self, message: bytes) -> None:
         async with self._lock:
-            self.ensure_open()
             await self._send_message(message)
 
     async def send_query(self, message: bytes, response_size: int) -> bytes:
         async with self._lock:
-            self.ensure_open()
             await self._send_message(message)
             return await self._receive_response(response_size)
 
-    async def close(self) -> None:
-        async with self._lock:
-            self.ensure_open()
-            self.stream.close()
-            self.stream = None
-
     async def _send_message(self, message):
-        await self.stream.write_async(message)
+        await self._stream.write_async(message)
 
     async def _receive_response(self, size):
-        return await self.stream.read_async(size)
+        return await self._stream.read_async(size)
+
+    async def close(self) -> None:
+        async with self._lock:
+            self._stream.close()
+            self.__stream = None

tests/backends/epics/test_gui.py

@@ -6,6 +6,7 @@
     SignalRW,
     SignalW,
     SignalX,
+    TextFormat,
     TextRead,
     TextWrite,
     ToggleButton,
@@ -53,7 +54,7 @@ def test_get_components(mapping):
         SignalRW(
             name="StringEnum",
             read_pv="DEVICE:StringEnum_RBV",
-            read_widget=TextRead(format="string"),
+            read_widget=TextRead(format=TextFormat.string),
             write_pv="DEVICE:StringEnum",
             write_widget=ComboBox(choices=["red", "green", "blue"]),
         ),

tests/backends/epics/test_ioc_system.py

@@ -1,10 +1,13 @@
+from p4p import Value
 from p4p.client.thread import Context
 
 
 def test_ioc(ioc: None):
     ctxt = Context("pva")
 
-    parent_pvi = ctxt.get("DEVICE:PVI").todict()
+    _parent_pvi = ctxt.get("DEVICE:PVI")
+    assert isinstance(_parent_pvi, Value)
+    parent_pvi = _parent_pvi.todict()
     assert all(f in parent_pvi for f in ("alarm", "display", "timeStamp", "value"))
     assert parent_pvi["display"] == {"description": "The records in this controller"}
     assert parent_pvi["value"] == {
@@ -14,7 +17,9 @@ def test_ioc(ioc: None):
     }
 
     child_pvi_pv = parent_pvi["value"]["child"]["d"]
-    child_pvi = ctxt.get(child_pvi_pv).todict()
+    _child_pvi = ctxt.get(child_pvi_pv)
+    assert isinstance(_child_pvi, Value)
+    child_pvi = _child_pvi.todict()
     assert all(f in child_pvi for f in ("alarm", "display", "timeStamp", "value"))
     assert child_pvi["display"] == {"description": "The records in this controller"}
     assert child_pvi["value"] == {

tests/conftest.py

@@ -106,7 +106,7 @@ def ioc():
 
     start_time = time.monotonic()
     while "iocRun: All initialization complete" not in (
-        process.stdout.readline().strip()
+        process.stdout.readline().strip()  # type: ignore
     ):
         if time.monotonic() - start_time > 10:
             raise TimeoutError("IOC did not start in time")