Merge pull request #13 from harp-tech/gl-dev

Refactor modules for type checking

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,37 +5,44 @@
 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 BaseModel, BeforeValidator, ConfigDict, Field, RootModel, conint, field_serializer
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    RootModel,
+    field_serializer,
+)
 
 
 class PayloadType(str, Enum):
-    U8 = 'uint8'
-    S8 = 'int8'
-    U16 = 'uint16'
-    S16 = 'int16'
-    U32 = 'uint32'
-    S32 = 'int32'
-    U64 = 'uint64'
-    S64 = 'int64'
-    Float = 'float32'
+    U8 = "uint8"
+    S8 = "int8"
+    U16 = "uint16"
+    S16 = "int16"
+    U32 = "uint32"
+    S32 = "int32"
+    U64 = "uint64"
+    S64 = "int64"
+    Float = "float32"
 
 
 class Access(Enum):
-    Read = 'Read'
-    Write = 'Write'
-    Event = 'Event'
+    Read = "Read"
+    Write = "Write"
+    Event = "Event"
 
 
 class MaskValueItem(BaseModel):
     model_config = ConfigDict(
-        extra='forbid',
+        extra="forbid",
     )
-    value: int = Field(..., description='Specifies the numerical mask value.')
+    value: int = Field(..., description="Specifies the numerical mask value.")
     description: Optional[str] = Field(
-        None, description='Specifies a summary description of the mask value function.'
+        None, description="Specifies a summary description of the mask value function."
     )
 
     def __int__(self):
@@ -48,70 +55,70 @@ class MaskValue(RootModel[Union[int, MaskValueItem]]):
 
 class BitMask(BaseModel):
     description: Optional[str] = Field(
-        None, description='Specifies a summary description of the bit mask function.'
+        None, description="Specifies a summary description of the bit mask function."
     )
     bits: Dict[str, MaskValue]
 
 
 class GroupMask(BaseModel):
     description: Optional[str] = Field(
-        None, description='Specifies a summary description of the group mask function.'
+        None, description="Specifies a summary description of the group mask function."
     )
     values: Dict[str, MaskValue]
 
 
 class MaskType(RootModel[str]):
     root: str = Field(
         ...,
-        description='Specifies the name of the bit mask or group mask used to represent the payload value.',
+        description="Specifies the name of the bit mask or group mask used to represent the payload value.",
     )
 
 
 class InterfaceType(RootModel[str]):
     root: str = Field(
         ...,
-        description='Specifies the name of the type used to represent the payload value in the high-level interface.',
+        description="Specifies the name of the type used to represent the payload value in the high-level interface.",
     )
 
 
 class Converter(Enum):
-    None_ = 'None'
-    Payload = 'Payload'
-    RawPayload = 'RawPayload'
+    None_ = "None"
+    Payload = "Payload"
+    RawPayload = "RawPayload"
 
 
 class MinValue(RootModel[float]):
     root: float = Field(
         ...,
-        description='Specifies the minimum allowable value for the payload or payload member.',
+        description="Specifies the minimum allowable value for the payload or payload member.",
     )
 
 
 class MaxValue(RootModel[float]):
     root: float = Field(
         ...,
-        description='Specifies the maximum allowable value for the payload or payload member.',
+        description="Specifies the maximum allowable value for the payload or payload member.",
     )
 
 
 class DefaultValue(RootModel[float]):
     root: float = Field(
         ...,
-        description='Specifies the default value for the payload or payload member.',
+        description="Specifies the default value for the payload or payload member.",
     )
 
 
 class PayloadMember(BaseModel):
     mask: Optional[int] = Field(
         None,
-        description='Specifies the mask used to read and write this payload member.',
+        description="Specifies the mask used to read and write this payload member.",
     )
     offset: Optional[int] = Field(
         None,
-        description='Specifies the payload array offset where this payload member is stored.',
+        description="Specifies the payload array offset where this payload member is stored.",
     )
     description: Optional[str] = Field(
-        None, description='Specifies a summary description of the payload member.'
+        None, description="Specifies a summary description of the payload member."
     )
     minValue: Optional[MinValue] = None
     maxValue: Optional[MaxValue] = None
@@ -122,68 +129,74 @@ class PayloadMember(BaseModel):
 
 
 class Visibility(Enum):
-    public = 'public'
-    private = 'private'
+    public = "public"
+    private = "private"
 
 
 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.'
+        ..., description="Specifies the expected use of the register."
     )
     description: Optional[str] = Field(
-        None, description='Specifies a summary description of the register function.'
+        None, description="Specifies a summary description of the register function."
     )
     minValue: Optional[MinValue] = None
     maxValue: Optional[MaxValue] = None
     defaultValue: Optional[DefaultValue] = None
     maskType: Optional[MaskType] = None
     visibility: Optional[Visibility] = Field(
         None,
-        description='Specifies whether the register function is exposed in the high-level interface.',
+        description="Specifies whether the register function is exposed in the high-level interface.",
     )
     volatile: Optional[bool] = Field(
         None,
-        description='Specifies whether register values can be saved in non-volatile memory.',
+        description="Specifies whether register values can be saved in non-volatile memory.",
     )
     payloadSpec: Optional[Dict[str, PayloadMember]] = None
     interfaceType: Optional[InterfaceType] = None
     converter: Optional[Converter] = None
 
-    @field_serializer('type')
+    @field_serializer("type")
     def _serialize_type(self, type: PayloadType):
         return type.name
 
 
 class Registers(BaseModel):
     registers: Dict[str, Register] = Field(
         ...,
-        description='Specifies the collection of registers implementing the device function.',
+        description="Specifies the collection of registers implementing the device function.",
     )
     bitMasks: Optional[Dict[str, BitMask]] = Field(
         None,
-        description='Specifies the collection of masks available to be used with the different registers.',
+        description="Specifies the collection of masks available to be used with the different registers.",
     )
     groupMasks: Optional[Dict[str, GroupMask]] = Field(
         None,
-        description='Specifies the collection of group masks available to be used with the different registers.',
+        description="Specifies the collection of group masks available to be used with the different registers.",
     )
 
 
 class Model(Registers):
-    device: str = Field(..., description='Specifies the name of the device.')
+    device: str = Field(..., description="Specifies the name of the device.")
     whoAmI: int = Field(
-        ..., description='Specifies the unique identifier for this device type.'
+        ..., description="Specifies the unique identifier for this device type."
     )
     firmwareVersion: str = Field(
-        ..., description='Specifies the semantic version of the device firmware.'
+        ..., description="Specifies the semantic version of the device firmware."
     )
     hardwareTargets: str = Field(
-        ..., description='Specifies the semantic version of the device hardware.'
+        ..., description="Specifies the semantic version of the device hardware."
     )

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)

pyproject.toml

@@ -65,7 +65,6 @@ extend-exclude = '''
 (
   ^/LICENSE
   ^/README.md
-  | harp/model.py
   | reflex-generator
 )
 '''