Refactor modules for typecheck support

Author: glopesdev

harp/io.py

@@ -84,7 +84,7 @@ def read(
         time = micros * _SECONDS_PER_TICK + seconds
         payloadtype = payloadtype & ~0x10
         if epoch is not None:
-            time = epoch + pd.to_timedelta(time, "s")
+            time = epoch + pd.to_timedelta(time, "s") # type: ignore
         index = pd.Series(time)
         index.name = "time"
 
@@ -111,6 +111,6 @@ def read(
         msgtype = np.ndarray(
             nrows, dtype=np.uint8, buffer=data, offset=0, strides=stride
         )
-        msgtype = pd.Categorical.from_codes(msgtype, categories=_messagetypes)
+        msgtype = pd.Categorical.from_codes(msgtype, categories=_messagetypes) # type: ignore
         result["type"] = msgtype
     return result

harp/model.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 from typing_extensions import Annotated
 
 from pydantic import (
@@ -14,7 +14,6 @@
     ConfigDict,
     Field,
     RootModel,
-    conint,
     field_serializer,
 )
 
@@ -135,13 +134,19 @@ class Visibility(Enum):
 
 
 class Register(BaseModel):
-    address: conint(le=255) = Field(
-        ..., description="Specifies the unique 8-bit address of the register."
-    )
+    address: Annotated[
+        int,
+        Field(
+            le=255, description="Specifies the unique 8-bit address of the register."
+        ),
+    ]
     type: Annotated[PayloadType, BeforeValidator(lambda v: PayloadType[v])]
-    length: Optional[conint(ge=1)] = Field(
-        default=1, description="Specifies the length of the register payload."
-    )
+    length: Annotated[
+        Optional[int],
+        Field(
+            ge=1, default=1, description="Specifies the length of the register payload."
+        ),
+    ]
     access: Union[Access, List[Access]] = Field(
         ..., description="Specifies the expected use of the register."
     )

harp/reader.py

@@ -2,18 +2,20 @@
 from math import log2
 from os import PathLike
 from functools import partial
+from numpy import dtype
 from pandas import DataFrame, Series
 from typing import Any, BinaryIO, Iterable, Callable, Optional, Union
 from pandas._typing import Axes
 from harp.model import BitMask, GroupMask, Model, PayloadMember, Register
 from harp.io import read
+from harp.schema import read_schema
 
 _camel_to_snake_regex = re.compile(r"(?<!^)(?=[A-Z])")
 
 
 class RegisterReader:
     register: Register
-    read: Callable[[str], DataFrame]
+    read: Callable[[Union[str, bytes, PathLike[Any], BinaryIO]], DataFrame]
 
     def __init__(
         self,
@@ -89,14 +91,14 @@ def _create_payloadmember_parser(device: Model, member: PayloadMember):
     if offset is None:
         offset = 0
 
-    shift = None
+    shift = 0
     if member.mask is not None:
         shift = _mask_shift(member.mask)
 
     lookup = None
     if member.maskType is not None:
         key = member.maskType.root
-        groupMask = device.groupMasks.get(key)
+        groupMask = None if device.groupMasks is None else device.groupMasks.get(key)
         if groupMask is not None:
             lookup = _create_groupmask_lookup(groupMask)
 
@@ -109,7 +111,7 @@ def parser(df: DataFrame):
         if member.mask is not None:
             series = series & member.mask
             if shift > 0:
-                series = Series(series.values >> shift, series.index)
+                series = Series(series.values >> shift, series.index) # type: ignore
         if is_boolean:
             series = series != 0
         elif lookup is not None:
@@ -126,7 +128,7 @@ def reader(
         data = read(
             file,
             address=register.address,
-            dtype=register.type,
+            dtype=dtype(register.type),
             length=register.length,
             columns=columns,
         )
@@ -141,13 +143,13 @@ def _create_register_parser(device: Model, name: str):
 
     if register.maskType is not None:
         key = register.maskType.root
-        bitMask = device.bitMasks.get(key)
+        bitMask = None if device.bitMasks is None else device.bitMasks.get(key)
         if bitMask is not None:
             parser = _create_bitmask_parser(bitMask)
             reader = _compose(parser, reader)
             return RegisterReader(register, reader)
 
-        groupMask = device.groupMasks.get(key)
+        groupMask = None if device.groupMasks is None else device.groupMasks.get(key)
         if groupMask is not None:
             parser = _create_groupmask_parser(name, groupMask)
             reader = _compose(parser, reader)
@@ -170,7 +172,8 @@ def parser(df: DataFrame):
     return RegisterReader(register, reader)
 
 
-def create_reader(device: Model):
+def create_reader(file: Union[str, PathLike], include_common_registers: bool = True):
+    device = read_schema(file, include_common_registers)
     reg_readers = {
         name: _create_register_parser(device, name) for name in device.registers.keys()
     }

harp/schema.py

@@ -8,24 +8,26 @@
 
 
 def _read_common_registers(file: Union[str, PathLike, TextIO]) -> Registers:
-    try:
+    if isinstance(file, TextIO):
+        return parse_yaml_raw_as(Registers, file.read())
+    else:
         with open(file) as fileIO:
             return _read_common_registers(fileIO)
-    except TypeError:
-        return parse_yaml_raw_as(Registers, file.read())
 
 
 def read_schema(
     file: Union[str, PathLike, TextIO], include_common_registers: bool = True
 ) -> Model:
-    try:
-        with open(file) as fileIO:
-            return read_schema(fileIO, include_common_registers)
-    except TypeError:
+    if isinstance(file, TextIO):
         schema = parse_yaml_raw_as(Model, file.read())
         if not "WhoAmI" in schema.registers and include_common_registers:
             common = _read_common_registers(_common_yaml_path)
             schema.registers = dict(common.registers, **schema.registers)
-            schema.bitMasks = dict(common.bitMasks, **schema.bitMasks)
-            schema.groupMasks = dict(common.groupMasks, **schema.groupMasks)
+            if schema.bitMasks is not None and common.bitMasks is not None:
+                schema.bitMasks = dict(common.bitMasks, **schema.bitMasks)
+            if schema.groupMasks is not None and common.groupMasks is not None:
+                schema.groupMasks = dict(common.groupMasks, **schema.groupMasks)
         return schema
+    else:
+        with open(file) as fileIO:
+            return read_schema(fileIO, include_common_registers)

tests/test_io.py

@@ -1,6 +1,7 @@
 import pytest
 import numpy as np
-from typing import Iterable, Optional, Type
+from os import PathLike
+from typing import Iterable, Optional, Type, Union
 from contextlib import nullcontext
 from pytest import mark
 from pathlib import Path
@@ -12,7 +13,7 @@
 
 @dataclass
 class DataFileParam:
-    path: str
+    path: Union[str, PathLike]
     expected_rows: int
     expected_cols: Optional[Iterable[str]] = None
     expected_address: Optional[int] = None
@@ -36,7 +37,7 @@ def __post_init__(self):
     DataFileParam(
         path="data/device_0.bin",
         expected_rows=1,
-        expected_dtype="uint8",  # actual dtype is uint16
+        expected_dtype=np.dtype("uint8"),  # actual dtype is uint16
         expected_error=ValueError,
     ),
     DataFileParam(
@@ -55,7 +56,7 @@ def __post_init__(self):
 @mark.parametrize("dataFile", testdata)
 def test_read(dataFile: DataFileParam):
     context = pytest.raises if dataFile.expected_error is not None else nullcontext
-    with context(dataFile.expected_error):
+    with context(dataFile.expected_error):  # type: ignore
         data = read(
             dataFile.path,
             address=dataFile.expected_address,