Merge pull request #82 from hephy-dd/type_hints

Adding missing type hints

changelog

@@ -7,11 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-## [1.1.0] - 2024-09-11
+## [1.1.0] - 2024-09-12
 
 ### Added
-- Additional HEPHY EnvironBox commands for driver and emulator (#74).
+- Missing type hints for static type checking (#82).
 - Generic base classes `PowerSupply` and `PowerSupplyChannel` and `LightSource`.
+- Additional HEPHY EnvironBox commands for driver and emulator (#74).
 - Rhode&Schwarz NGE100 power supply (#77).
 - Photonic F3000 LED light source (#75).
 - Thorlabs PM100 USB optical power meter (#78).

src/comet/driver/__init__.py

@@ -1 +1,3 @@
 from .driver import Driver
+
+__all__ = ["Driver"]

src/comet/driver/corvus/__init__.py

@@ -1 +1,3 @@
 from .venus1 import Venus1
+
+__all__ = ["Venus1"]

src/comet/driver/cts/__init__.py

@@ -1 +1,3 @@
 from .itc import ITC
+
+__all__ = ["ITC"]

src/comet/driver/cts/itc.py

@@ -1,6 +1,6 @@
 import datetime
-from typing import Tuple
 from collections import namedtuple
+from typing import Union
 
 from comet.driver import Driver
 
@@ -10,7 +10,7 @@
 class ITCDriver(Driver):
     """ITC driver base class."""
 
-    def query_bytes(self, message, count) -> str:
+    def query_bytes(self, message: Union[str, bytes], count: int) -> str:
         """Raw query for bytes.
 
         >>> instr.query_bytes("P", 4)
@@ -23,8 +23,7 @@ def query_bytes(self, message, count) -> str:
 
 
 class AnalogChannel(ITCDriver):
-
-    CHANNELS = {
+    CHANNELS: dict[int, bytes] = {
         1: b"A0",
         2: b"A1",
         3: b"A2",
@@ -45,7 +44,7 @@ class AnalogChannel(ITCDriver):
     """Mapping analog channel index to channel code. When writing convert to
     lower case."""
 
-    def __getitem__(self, index: int) -> Tuple[float, float]:
+    def __getitem__(self, index: int) -> tuple[float, float]:
         """Read analog channel, returns tuple containing actual value and target value.
 
         >>> instr.analog_channel[1] # read temperature target/actual
@@ -57,7 +56,7 @@ def __getitem__(self, index: int) -> Tuple[float, float]:
             raise RuntimeError(f"invalid channel returned: '{result}'")
         return float(actual), float(target)
 
-    def __setitem__(self, index: int, value: float):
+    def __setitem__(self, index: int, value: float) -> None:
         """Set target value for analog channel.
 
         >>> instr.set_analog_channel[1] = 42.0
@@ -73,7 +72,7 @@ def __setitem__(self, index: int, value: float):
 class ITC(ITCDriver):
     """Interface for CTS Climate Chambers."""
 
-    WARNING_MESSAGES = {
+    WARNING_MESSAGES: dict[str, str] = {
         "\x01": "Wassernachfüllen",
         "\x02": "Temp. Toleranzband Oben",
         "\x03": "Temp. Toleranzband Unten",
@@ -83,7 +82,7 @@ class ITC(ITCDriver):
     }
     """Warning messages."""
 
-    ERROR_MESSAGES = {
+    ERROR_MESSAGES: dict[str, str] = {
         "\x31": "Temperatur Grenze Min 08-B1",
         "\x32": "Temperatur Grenze Max 08-B1",
         "\x33": "Temp. Begrenzer Pruefr. 01-F1.1",
@@ -114,11 +113,11 @@ class ITC(ITCDriver):
     Status = namedtuple("Status", ("running", "warning", "error", "channels"))
     """Status type container."""
 
-    def __init__(self, resource):
+    def __init__(self, resource) -> None:
         super().__init__(resource)
         self.analog_channel = AnalogChannel(resource)
 
-    def identify(self):
+    def identify(self) -> str:
         """Returns instrument identification."""
         self.time  # perform device access
         return "ITC climate chamber"
@@ -134,7 +133,7 @@ def time(self) -> datetime.datetime:
         return datetime.datetime.strptime(result, "T%d%m%y%H%M%S")
 
     @time.setter
-    def time(self, dt: datetime.datetime):
+    def time(self, dt: datetime.datetime) -> None:
         """Update device date and time, returns updated data and time as datetime object.
 
         >>> instr.time = datetime.datetime.now()
@@ -145,7 +144,7 @@ def time(self, dt: datetime.datetime):
             raise RuntimeError("failed to set date and time")
 
     @property
-    def status(self) -> object:
+    def status(self) -> Status:
         """Returns device status as object.
 
         >>> instr.status
@@ -160,7 +159,7 @@ def status(self) -> object:
         error_nr = result[9]
         warning = type(self).WARNING_MESSAGES[error_nr] if is_error and error_nr in type(self).WARNING_MESSAGES else None
         error = type(self).ERROR_MESSAGES[error_nr] if is_error and error_nr in type(self).ERROR_MESSAGES else None
-        return self.Status(running, warning, error, channels)
+        return type(self).Status(running, warning, error, channels)
 
     @property
     def error_message(self) -> str:
@@ -183,7 +182,7 @@ def program(self) -> int:
         return int(result[1:])
 
     @program.setter
-    def program(self, number: int):
+    def program(self, number: int) -> None:
         """Starts a program. Returns program number or 0 for no program.
 
         >>> instr.program = 42

src/comet/driver/driver.py

@@ -1,7 +1,8 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 
 
 class Driver(ABC):
+    """Base class for instrument drivers."""
 
-    def __init__(self, resource):
+    def __init__(self, resource) -> None:
         self.resource = resource

src/comet/driver/generic/switching_matrix.py

@@ -1,23 +1,22 @@
 from abc import abstractmethod
-from typing import List
 
 from .instrument import Instrument
 
 __all__ = ["SwitchingMatrix"]
 
 
 class SwitchingMatrix(Instrument):
-    CHANNELS: List[str] = []
+    CHANNELS: list[str] = []
 
     @property
     @abstractmethod
-    def closed_channels(self) -> List[str]: ...
+    def closed_channels(self) -> list[str]: ...
 
     @abstractmethod
-    def close_channels(self, channels: List[str]) -> None: ...
+    def close_channels(self, channels: list[str]) -> None: ...
 
     @abstractmethod
-    def open_channels(self, channels: List[str]) -> None: ...
+    def open_channels(self, channels: list[str]) -> None: ...
 
     @abstractmethod
     def open_all_channels(self) -> None: ...

src/comet/driver/hephy/__init__.py

@@ -1,3 +1,5 @@
 from .brandbox import BrandBox
 from .environbox import EnvironBox
 from .shuntbox import ShuntBox
+
+__all__ = ["BrandBox", "EnvironBox", "ShuntBox"]

src/comet/driver/hephy/brandbox.py

@@ -1,27 +1,25 @@
 import re
-from typing import Dict, List, Optional, Tuple
+from typing import Optional
 
 from comet.driver.generic import InstrumentError
 from comet.driver.generic.switching_matrix import SwitchingMatrix
 from comet.utils import combine_matrix
 
-__all__ = ['BrandBox']
+__all__ = ["BrandBox"]
 
-ERROR_MESSAGES: Dict[int, str] = {
-    99: "Invalid command"
-}
+ERROR_MESSAGES: dict[int, str] = {99: "Invalid command"}
 
 
-def split_channels(channels: str) -> List[str]:
-    return [channel.strip() for channel in channels.split(',') if channel.strip()]
+def split_channels(channels: str) -> list[str]:
+    return [channel.strip() for channel in channels.split(",") if channel.strip()]
 
 
-def join_channels(channels: List[str]) -> str:
-    return ','.join([format(channel).strip() for channel in channels])
+def join_channels(channels: list[str]) -> str:
+    return ",".join([format(channel).strip() for channel in channels])
 
 
 def parse_error(response: str) -> Optional[InstrumentError]:
-    m = re.match(r'^err(\d+)', response.lower())
+    m = re.match(r"^err(\d+)", response.lower())
     if m:
         code = int(m.group(1))
         message = ERROR_MESSAGES.get(code, "")
@@ -30,21 +28,20 @@ def parse_error(response: str) -> Optional[InstrumentError]:
 
 
 class BrandBox(SwitchingMatrix):
+    CHANNELS: list[str] = combine_matrix("ABC", "12")
 
-    CHANNELS = combine_matrix('ABC', '12')
-
-    _error_queue: List[InstrumentError] = []
+    _error_queue: list[InstrumentError] = []
 
     def identify(self) -> str:
-        return self.query('*IDN?')
+        return self.query("*IDN?")
 
     def reset(self) -> None:
         self._error_queue.clear()
-        self.write('*RST')
+        self.write("*RST")
 
     def clear(self) -> None:
         self._error_queue.clear()
-        self.write('*CLS')
+        self.write("*CLS")
 
     # Error queue
 
@@ -56,21 +53,21 @@ def next_error(self) -> Optional[InstrumentError]:
     # Switching matrix
 
     @property
-    def closed_channels(self) -> List[str]:
-        channels = self.query(':CLOS:STAT?')
+    def closed_channels(self) -> list[str]:
+        channels = self.query(":CLOS:STAT?")
         return split_channels(channels)
 
-    def close_channels(self, channels: List[str]) -> None:
+    def close_channels(self, channels: list[str]) -> None:
         channel_list = join_channels(sorted(channels))
-        self.write(f':CLOS {channel_list}')
+        self.write(f":CLOS {channel_list}")
 
-    def open_channels(self, channels: List[str]) -> None:
+    def open_channels(self, channels: list[str]) -> None:
         channel_list = join_channels(sorted(channels))
-        self.write(f':OPEN {channel_list}')
+        self.write(f":OPEN {channel_list}")
 
     def open_all_channels(self) -> None:
         channel_list = join_channels(type(self).CHANNELS)
-        self.write(f':OPEN {channel_list}')
+        self.write(f":OPEN {channel_list}")
 
     # Helper
 
@@ -79,7 +76,7 @@ def query(self, message: str) -> str:
         error = parse_error(response)
         if error:
             self._error_queue.append(error)
-            return ''
+            return ""
         return response
 
     def write(self, message: str) -> None:

src/comet/driver/hephy/environbox.py

@@ -1,12 +1,12 @@
 import re
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 
 from comet.driver.generic import Instrument
 from comet.driver.generic import InstrumentError
 
 __all__ = ["EnvironBox"]
 
-ERROR_MESSAGES: Dict[int, str] = {
+ERROR_MESSAGES: dict[int, str] = {
     1: "RTC not running",
     2: "RTC read error",
     80: "DAC not found",
@@ -32,7 +32,7 @@ def parse_error(response: str) -> Optional[InstrumentError]:
     return None
 
 
-def parse_pc_data(response: str) -> Dict[str, Any]:
+def parse_pc_data(response: str) -> dict[str, Any]:
     values = response.split(",")
     relay_status = int(values[23])
     return {
@@ -85,8 +85,7 @@ def parse_pc_data(response: str) -> Dict[str, Any]:
 
 
 class EnvironBox(Instrument):
-
-    _error_queue: List[InstrumentError] = []
+    _error_queue: list[InstrumentError] = []
 
     def identify(self) -> str:
         return self.query("*IDN?")
@@ -225,7 +224,7 @@ def set_door_auto_light(self, state: bool) -> None:
         value = {self.DOOR_AUTO_LIGHT_OFF: "OFF", self.DOOR_AUTO_LIGHT_ON: "ON"}[state]
         self.write(f"SET:DOOR_AUTO_LIGHT {value}")
 
-    def get_data(self) -> Dict[str, Any]:
+    def get_data(self) -> dict[str, Any]:
         """Return dictionary of PC_DATA."""
         return parse_pc_data(self.query("GET:PC_DATA ?"))
 

src/comet/driver/hephy/shuntbox.py

@@ -1,18 +1,16 @@
 import re
-from typing import Dict, List, Optional, Tuple
+from typing import Optional
 
 from comet.driver.generic import Instrument
 from comet.driver.generic import InstrumentError
 
-__all__ = ['ShuntBox']
+__all__ = ["ShuntBox"]
 
-ERROR_MESSAGES: Dict[int, str] = {
-    99: "Unknown command"
-}
+ERROR_MESSAGES: dict[int, str] = {99: "Unknown command"}
 
 
 def parse_error(response: str) -> Optional[InstrumentError]:
-    m = re.match(r'^err(\d+)', response.lower())
+    m = re.match(r"^err(\d+)", response.lower())
     if m:
         code = int(m.group(1))
         message = ERROR_MESSAGES.get(code, "")
@@ -21,15 +19,14 @@ def parse_error(response: str) -> Optional[InstrumentError]:
 
 
 class ShuntBox(Instrument):
-
-    _error_queue: List[InstrumentError] = []
+    _error_queue: list[InstrumentError] = []
 
     def identify(self) -> str:
-        return self.query('*IDN?')
+        return self.query("*IDN?")
 
     def reset(self) -> None:
         self._error_queue.clear()
-        self.write('*RST')
+        self.write("*RST")
 
     def clear(self) -> None:
         self._error_queue.clear()
@@ -49,7 +46,7 @@ def query(self, message: str) -> str:
         error = parse_error(response)
         if error:
             self._error_queue.append(error)
-            return ''
+            return ""
         return response
 
     def write(self, message: str) -> None:

src/comet/driver/itk/__init__.py

@@ -1,2 +1,4 @@
 from .corvustt import CorvusTT
 from .hydra import Hydra
+
+__all__ = ["CorvusTT", "Hydra"]