Skip to content

Commit

Permalink
add support for numpy array writes
Browse files Browse the repository at this point in the history
  • Loading branch information
WatcherBox committed Feb 28, 2024
1 parent 2a78a49 commit f340089
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 7 deletions.
1 change: 1 addition & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ packages = find:
install_requires =
PyYAML>=5.3.1
pyserial>=3.4
numpy>=1.19.4

[options.entry_points]
console_scripts =
Expand Down
29 changes: 27 additions & 2 deletions simple_rpc/io.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import numpy as np
from typing import Any, BinaryIO
from struct import calcsize, pack, unpack

Expand Down Expand Up @@ -43,9 +44,12 @@ def _write_basic(
if basic_type == 's':
stream.write(value + b'\0')
return
elif isinstance(basic_type, np.ndarray):
print(f"writing array: {basic_type.itemsize}, {value.size}")
stream.write(value.tobytes())
return

full_type = (endianness + basic_type).encode('utf-8')

stream.write(pack(full_type, cast(basic_type)(value)))


Expand Down Expand Up @@ -77,14 +81,29 @@ def read(
:returns: Object of type {obj_type}.
"""

print(f"reading: {obj_type}")

if isinstance(obj_type, np.ndarray):
# print(f"reading array: {size_t}, {obj_type}, {obj_type.itemsize}")
length = _read_basic(stream, endianness, size_t)
# print(f"length: {length}")
return np.frombuffer(
stream.read(length * obj_type.itemsize), obj_type.dtype)

if isinstance(obj_type, list):
# print(f"reading list: {obj_type}")
length = _read_basic(stream, endianness, size_t)
# print(f"length: {length}")

return [
read(stream, endianness, size_t, item) for _ in range(length)
for item in obj_type]

if isinstance(obj_type, tuple):
return tuple(
read(stream, endianness, size_t, item) for item in obj_type)

return _read_basic(stream, endianness, obj_type)


Expand All @@ -103,9 +122,15 @@ def write(
:arg obj_type: Type object.
:arg obj: Object of type {obj_type}.
"""

if isinstance(obj_type, list):
# print(f" size_t: {size_t}, len:{len(obj) // len(obj_type)}")
_write_basic(stream, endianness, size_t, len(obj) // len(obj_type))
if isinstance(obj_type, list) or isinstance(obj_type, tuple):
if isinstance(obj_type, np.ndarray):
# print(f"writing array: {size_t}, {obj.size}, {obj.dtype}, obj_tpye: {obj_type}")
_write_basic(stream, endianness, size_t, obj.size)
_write_basic(stream, endianness, obj_type, obj)
elif isinstance(obj_type, list) or isinstance(obj_type, tuple):
for item_type, item in zip(obj_type * len(obj), obj):
write(stream, endianness, size_t, item_type, item)
else:
Expand Down
41 changes: 37 additions & 4 deletions simple_rpc/protocol.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,25 @@
from typing import Any, BinaryIO

import numpy as np
from .io import cast, read_byte_string

from itertools import chain


dtype_map = {
'b': np.int8,
'B': np.uint8,
'h': np.int16,
'H': np.uint16,
'i': np.int32,
'I': np.uint32,
'l': np.int32,
'L': np.uint32,
'q': np.int64,
'Q': np.uint64,
'f': np.float32,
'd': np.float64,
'?': np.bool_,
'c': np.byte # Note: 'c' in struct is a single byte; for strings, consider np.bytes_ or np.chararray.
}

def _parse_type(type_str: bytes) -> Any:
"""Parse a type definition string.
Expand All @@ -15,14 +33,20 @@ def _construct_type(tokens: tuple):

for token in tokens:
if token == b'[':
obj_type.append(_construct_type(tokens))
next_token = next(tokens, None)
if next_token not in (b'[', b'(') and next_token is not None:
dtype = _get_dtype(next_token)
obj_type.append(np.array([], dtype=dtype))
assert next(tokens, None) == b']', "Expected closing bracket"
else:
tokens = chain([next_token], tokens)
obj_type.extend(_construct_type(tokens))
elif token == b'(':
obj_type.append(tuple(_construct_type(tokens)))
elif token in (b')', b']'):
break
else:
obj_type.append(token.decode())

return obj_type

obj_type = _construct_type((bytes([char]) for char in type_str))
Expand All @@ -33,6 +57,15 @@ def _construct_type(tokens: tuple):
return ''
return obj_type[0]

def _get_dtype(type_str: bytes) -> Any:
"""Get the NumPy data type of a type definition string.
:arg type_str: Type definition string.
:returns: NumPy data type.
"""
return dtype_map.get(type_str, np.byte)


def _type_name(obj_type: Any) -> str:
"""Python type name of a C object type.
Expand Down
9 changes: 8 additions & 1 deletion simple_rpc/simple_rpc.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from functools import wraps
from time import sleep
import numpy as np
from types import MethodType
from typing import Any, TextIO

Expand Down Expand Up @@ -86,6 +87,7 @@ def _write(self: object, obj_type: Any, obj: Any) -> None:
:arg obj_type: Type of the parameter.
:arg obj: Value of the parameter.
"""
# print(f"write obj_type: {obj_type}")
write(
self._connection, self.device['endianness'], self.device['size_t'],
obj_type, obj)
Expand All @@ -100,6 +102,7 @@ def _read(self: object, obj_type: Any) -> Any:
:returns: Return value.
"""
# print(f"read obj_type: {obj_type}")
return read(
self._connection, self.device['endianness'], self.device['size_t'],
obj_type)
Expand Down Expand Up @@ -180,11 +183,15 @@ def call_method(self: object, name: str, *args: Any) -> Any:

# Provide parameters (if any).
if method['parameters']:
# print(f"method['name']: {method['name']}")
# print(f"method['parameters']: {method['parameters']}")
for index, parameter in enumerate(method['parameters']):
self._write(parameter['fmt'], args[index])

# Read return value (if any).
if method['return']['fmt']:
if method['return']['fmt'] or isinstance(method['return']['fmt'], np.ndarray) :
# print(f"method['name']: {method['name']}")
# print(f"method['return']['fmt']: {method['return']['fmt']}")
return self._read(method['return']['fmt'])
return None

Expand Down

0 comments on commit f340089

Please sign in to comment.