diff --git a/setup.cfg b/setup.cfg index f4d4663..6b7778b 100644 --- a/setup.cfg +++ b/setup.cfg @@ -21,6 +21,7 @@ packages = find: install_requires = PyYAML>=5.3.1 pyserial>=3.4 + numpy>=1.19.4 [options.entry_points] console_scripts = diff --git a/simple_rpc/io.py b/simple_rpc/io.py index 9718a7b..fcd4a52 100644 --- a/simple_rpc/io.py +++ b/simple_rpc/io.py @@ -1,3 +1,4 @@ +import numpy as np from typing import Any, BinaryIO from struct import calcsize, pack, unpack @@ -43,9 +44,12 @@ def _write_basic( if basic_type == 's': stream.write(value + b'\0') return + + elif isinstance(value, np.ndarray): + stream.write(value.tobytes()) + return full_type = (endianness + basic_type).encode('utf-8') - stream.write(pack(full_type, cast(basic_type)(value))) @@ -82,9 +86,15 @@ def read( 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) + + if isinstance(obj_type, np.ndarray): + length = _read_basic(stream, endianness, size_t) + return np.frombuffer( + stream.read(length * obj_type.itemsize), obj_type.dtype) return _read_basic(stream, endianness, obj_type) @@ -104,14 +114,21 @@ def write( :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: _write_basic(stream, endianness, obj_type, obj) + + def until( condition: callable, f: callable, *args: Any, **kwargs: Any) -> None: """Call {f(*args, **kwargs)} until {condition} is true. diff --git a/simple_rpc/protocol.py b/simple_rpc/protocol.py index 30ed264..e368cff 100644 --- a/simple_rpc/protocol.py +++ b/simple_rpc/protocol.py @@ -1,7 +1,25 @@ +import numpy as np + from typing import Any, BinaryIO from .io import cast, read_byte_string +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. @@ -18,7 +36,12 @@ def _construct_type(tokens: tuple): obj_type.append(_construct_type(tokens)) elif token == b'(': obj_type.append(tuple(_construct_type(tokens))) - elif token in (b')', b']'): + elif token == b'{': + t = _construct_type(tokens) + assert len(t) == 1, 'only atomic types allowed in np arrays' + dtype = _get_dtype(t[0]) + obj_type.append(np.ndarray(dtype=dtype, shape=(1, 1))) + elif token in (b')', b']', b'}'): break else: obj_type.append(token.decode()) @@ -33,6 +56,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. @@ -41,6 +73,8 @@ def _type_name(obj_type: Any) -> str: :returns: Python type name. """ + if isinstance(obj_type, np.ndarray): + return '{' + ', '.join([_type_name(item) for item in obj_type]) + '}' if not obj_type: return '' if isinstance(obj_type, list): diff --git a/simple_rpc/simple_rpc.py b/simple_rpc/simple_rpc.py index e42a7d3..56905de 100644 --- a/simple_rpc/simple_rpc.py +++ b/simple_rpc/simple_rpc.py @@ -1,3 +1,6 @@ + +import numpy as np + from functools import wraps from time import sleep from types import MethodType @@ -184,7 +187,7 @@ def call_method(self: object, name: str, *args: Any) -> Any: 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): return self._read(method['return']['fmt']) return None