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

Add support for keeping message type

Author: glopesdev

harp/io.py

@@ -1,11 +1,21 @@
+from enum import IntEnum
 from os import PathLike
 from typing import Any, BinaryIO, Optional, Union
 from pandas._typing import Axes
 import numpy as np
 import pandas as pd
 
+
+class MessageType(IntEnum):
+    NA = 0
+    READ = 1
+    WRITE = 2
+    EVENT = 3
+
+
 _SECONDS_PER_TICK = 32e-6
-payloadtypes = {
+_messagetypes = [type.name for type in MessageType]
+_payloadtypes = {
     1: np.dtype(np.uint8),
     2: np.dtype(np.uint16),
     4: np.dtype(np.uint32),
@@ -24,6 +34,7 @@ def read(
     dtype: Optional[np.dtype] = None,
     length: Optional[int] = None,
     columns: Optional[Axes] = None,
+    keep_type: bool = False,
 ):
     """
     Read single-register Harp data from the specified file.
@@ -37,6 +48,7 @@ def read(
     :param length: Expected number of elements in register payload. If specified,
       the payload length of the first message in the file is used for validation.
     :param columns: The optional column labels to use for the data values.
+    :param keep_type: Specifies whether to include a column with the message type.
     :return: A pandas data frame containing message data, sorted by time.
     """
     data = np.fromfile(file, dtype=np.uint8)
@@ -68,7 +80,7 @@ def read(
         index.name = "time"
 
     payloadsize = stride - payloadoffset - 1
-    payloadtype = payloadtypes[payloadtype]
+    payloadtype = _payloadtypes[payloadtype]
     if dtype is not None and dtype != payloadtype:
         raise ValueError(f"expected payload type {dtype} but got {payloadtype}")
 
@@ -85,4 +97,11 @@ def read(
         strides=(stride, elementsize),
     )
 
-    return pd.DataFrame(payload, index=index, columns=columns)
+    result = pd.DataFrame(payload, index=index, columns=columns)
+    if keep_type:
+        msgtype = np.ndarray(
+            nrows, dtype=np.uint8, buffer=data, offset=0, strides=stride
+        )
+        msgtype = pd.Categorical.from_codes(msgtype, categories=_messagetypes)
+        result["type"] = msgtype
+    return result

tests/test_io.py

@@ -1,6 +1,6 @@
 import pytest
 import numpy as np
-from typing import Optional, Type
+from typing import Iterable, Optional, Type
 from contextlib import nullcontext
 from pytest import mark
 from pathlib import Path
@@ -14,10 +14,12 @@
 class DataFileParam:
     path: str
     expected_rows: int
+    expected_cols: Optional[Iterable[str]] = None
     expected_address: Optional[int] = None
     expected_dtype: Optional[np.dtype] = None
     expected_length: Optional[int] = None
     expected_error: Optional[Type[BaseException]] = None
+    keep_type: bool = False
 
     def __post_init__(self):
         self.path = datapath / self.path
@@ -43,10 +45,10 @@ def __post_init__(self):
         expected_length=2,  # actual length is 1
         expected_error=ValueError,
     ),
+    DataFileParam(path="data/write_0.bin", expected_address=0, expected_rows=4),
     DataFileParam(
-        path="data/write_0.bin",
-        expected_address=0,
-        expected_rows=4)
+        path="data/write_0.bin", expected_address=0, expected_rows=4, keep_type=True
+    ),
 ]
 
 
@@ -59,8 +61,15 @@ def test_read(dataFile: DataFileParam):
             address=dataFile.expected_address,
             dtype=dataFile.expected_dtype,
             length=dataFile.expected_length,
+            keep_type=dataFile.keep_type,
         )
         assert len(data) == dataFile.expected_rows
+        if dataFile.keep_type:
+            assert "type" in data.columns and data["type"].dtype == "category"
+
+        if dataFile.expected_cols is not None:
+            for col in dataFile.expected_cols:
+                assert col in data.columns
 
 
 if __name__ == "__main__":