From e09392e3c2376a8f25f959b980483639d3b729dc Mon Sep 17 00:00:00 2001 From: Alexey Timin Date: Fri, 17 Nov 2023 16:51:16 +0100 Subject: [PATCH 1/2] implement easy protobuf python module --- .gitignore | 4 +- .pre-commit-config.yaml | 1 - cpp/test_package/conanfile.py | 4 +- python/pkg/drift_protocol/easy/__init__.py | 3 + python/pkg/drift_protocol/easy/helpers.py | 0 python/pkg/drift_protocol/easy/package.py | 140 +++++++++++++++++++++ python/pkg/drift_protocol/easy/trigger.py | 53 ++++++++ python/setup.py | 2 +- python/tests/easy/test_package.py | 37 ++++++ python/tests/easy/test_trigger.py | 24 ++++ 10 files changed, 263 insertions(+), 5 deletions(-) create mode 100644 python/pkg/drift_protocol/easy/__init__.py create mode 100644 python/pkg/drift_protocol/easy/helpers.py create mode 100644 python/pkg/drift_protocol/easy/package.py create mode 100644 python/pkg/drift_protocol/easy/trigger.py create mode 100644 python/tests/easy/test_package.py create mode 100644 python/tests/easy/test_trigger.py diff --git a/.gitignore b/.gitignore index c9a3374..73a692a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,5 +3,7 @@ lib/ build/ dist/ python/pkg/drift_protocol.egg-info/ -python/pkg/drift_protocol +!python/pkg/drift_protocol/ +python/pkg/drift_protocol/* +!python/pkg/drift_protocol/easy/ !python/pkg/drift_protocol/__init__.py.in diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 91610fb..a80d1e7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,7 +16,6 @@ repos: name: black entry: black language: python - language_version: python3.8 args: - . - --check diff --git a/cpp/test_package/conanfile.py b/cpp/test_package/conanfile.py index d1422f3..5b19d98 100644 --- a/cpp/test_package/conanfile.py +++ b/cpp/test_package/conanfile.py @@ -19,12 +19,12 @@ def requirements(self): def build(self): cmake = CMake(self) cmake.configure() - cmake.build() + cmake.build_from_msg() def layout(self): cmake_layout(self) def test(self): if can_run(self): - cmd = os.path.join(self.cpp.build.bindirs[0], "test_package") + cmd = os.path.join(self.cpp.build_from_msg.bindirs[0], "test_package") self.run(cmd, env="conanrun") diff --git a/python/pkg/drift_protocol/easy/__init__.py b/python/pkg/drift_protocol/easy/__init__.py new file mode 100644 index 0000000..323514e --- /dev/null +++ b/python/pkg/drift_protocol/easy/__init__.py @@ -0,0 +1,3 @@ +from drift_protocol.easy.package import DriftPackage, StatusCode + +from drift_protocol.easy.trigger import TriggerMessage diff --git a/python/pkg/drift_protocol/easy/helpers.py b/python/pkg/drift_protocol/easy/helpers.py new file mode 100644 index 0000000..e69de29 diff --git a/python/pkg/drift_protocol/easy/package.py b/python/pkg/drift_protocol/easy/package.py new file mode 100644 index 0000000..f084a4b --- /dev/null +++ b/python/pkg/drift_protocol/easy/package.py @@ -0,0 +1,140 @@ +"""Drfit Package""" +import time +from typing import Optional, Dict + +from google.protobuf.any_pb2 import Any +from drift_protocol.common import DriftPackage as ProtoDriftPackage, StatusCode, DataPayload +from drift_bytes import Variant, OutputBuffer + +from drift_protocol.meta import TypedDataInfo, MetaInfo + + +class DriftPackage: + + def __init__(self): + self._proto = ProtoDriftPackage() + + @staticmethod + def parse(message): + """ + Parse drift package + + Args: + message: Message as bytes + + Returns: + Drift package + """ + package = DriftPackage() + package._proto.ParseFromString(message) + return package + + @staticmethod + def build_from_msg(message, pkg_id: Optional[int | float] = None, status: StatusCode = StatusCode.GOOD, **kwargs) -> "DriftPackage": + """ + Build drift package + + Args: + message: Message to be packed + pkg_id: ID as timestamp as float in seconds or int in milliseconds. If None, current time is used. + status: Status code + + Keyword Args: + publish_timestamp (int | float): Publish timestamp as float in seconds or int in milliseconds. + If None, current time is used. + source_timestamp (int | float): Source timestamp as float in seconds or int in milliseconds. + If None, current time is used. + Returns: + Drift package + """ + package = DriftPackage() + if pkg_id is None: + pkg_id = time.time() + + pkg_id = int(pkg_id * 1000) if isinstance(pkg_id, float) else pkg_id + package._proto.id = pkg_id + package._proto.status = status + + publish_timestamp = kwargs.get("publish_timestamp", time.time()) + if isinstance(publish_timestamp, float): + publish_timestamp = int(publish_timestamp * 1000) + package._proto.publish_timestamp.FromMilliseconds(publish_timestamp) + + source_timestamp = kwargs.get("source_timestamp", time.time()) + if isinstance(source_timestamp, float): + source_timestamp = int(source_timestamp * 1000) + package._proto.source_timestamp.FromMilliseconds(source_timestamp) + + msg = Any() + + if hasattr(message, "_proto"): + msg.Pack(message._proto) + else: + msg.Pack(message) + + package._proto.data.append(msg) + + return package + + @staticmethod + def build_from_typed_data(values: Dict[str, Variant.SUPPORTED_TYPES], + statuses: Optional[Dict[str, StatusCode]] = None, + pkg_id: Optional[int | float] = None, status: StatusCode = StatusCode.GOOD, **kwargs) -> "DriftPackage": + + statuses = statuses or {} + + info = TypedDataInfo() + + buffer = OutputBuffer(len(values)) + for idx, entry in enumerate(values.items()): + name, value = entry + buffer[idx] = value + + item = TypedDataInfo.Item() + item.name = name + item.type = statuses.get(name, StatusCode.GOOD) + info.items.append(item) + + data_payload = DataPayload() + data_payload.data = buffer.bytes() + + package = DriftPackage.build_from_msg(data_payload, pkg_id=pkg_id, status=status, **kwargs) + package._proto.meta.type = MetaInfo.TYPED_DATA + package._proto.meta.typed_data_info.CopyFrom(info) + + @property + def id(self) -> int: + """ + Package ID + """ + return self._proto.id + + @property + def status(self) -> StatusCode: + """ + Package status + """ + return self._proto.status + + @property + def publish_timestamp(self) -> float: + """ + Publish timestamp in seconds + """ + return self._proto.publish_timestamp.ToMilliseconds() / 1000 + + @property + def source_timestamp(self) -> float: + """ + Source timestamp in seconds + """ + return self._proto.source_timestamp.ToMilliseconds() / 1000 + + def unpack(self, msg): + """ + Unpack message it must be either a protobuf message or an "easy" message + """ + if hasattr(msg, "_proto"): + self._proto.data[0].Unpack(msg._proto) + else: + self._proto.data[0].Unpack(msg) diff --git a/python/pkg/drift_protocol/easy/trigger.py b/python/pkg/drift_protocol/easy/trigger.py new file mode 100644 index 0000000..e3f12f1 --- /dev/null +++ b/python/pkg/drift_protocol/easy/trigger.py @@ -0,0 +1,53 @@ +"""Trigger messages""" +import time +from typing import Dict, Optional + +from drift_protocol.trigger_service.trigger_message_pb2 import TriggerMessage as ProtoTriggerMessage + + +class TriggerMessage: + """Trigger message base class""" + + def __init__(self) -> None: + self._proto = ProtoTriggerMessage() + + @staticmethod + def parse(message: bytes) -> "TriggerMessage": + """ + Parse trigger message + + Args: + message: Message as bytes + + Returns: + Trigger message + """ + trigger = TriggerMessage() + trigger._proto.ParseFromString(message) + return trigger + + @staticmethod + def build(timestamp: Optional[int | float] = None) -> "TriggerMessage": + """ + Build trigger message + + Args: + timestamp: Timestamp as float in seconds or int in milliseconds. If None, current time is used. + + Returns: + Trigger message + """ + trigger = TriggerMessage() + if timestamp is None: + timestamp = time.time() + + timestamp = int(timestamp * 1000) if isinstance(timestamp, float) else timestamp + trigger._proto.timestamp.FromMilliseconds(timestamp) + return trigger + + @property + def timestamp(self) -> float: + """ + Timestamp in seconds + """ + return self._proto.timestamp.ToMilliseconds() / 1000 diff --git a/python/setup.py b/python/setup.py index c6625f6..cea6288 100644 --- a/python/setup.py +++ b/python/setup.py @@ -144,5 +144,5 @@ def packages(self): package_dir={"": "pkg"}, packages=LazyPackageFinder(finder=partial(find_packages, where="pkg")), python_requires=">=3.7", - install_requires=[f"protobuf>={PROTOBUF_VERSION}, <=3.20.3"], + install_requires=[f"protobuf>={PROTOBUF_VERSION}, <=3.20.3", "drift-bytes>=0.3.0"], ) diff --git a/python/tests/easy/test_package.py b/python/tests/easy/test_package.py new file mode 100644 index 0000000..a0e1d74 --- /dev/null +++ b/python/tests/easy/test_package.py @@ -0,0 +1,37 @@ +"""Unit test for package.py""" + +from drift_protocol.common import DriftPackage as ProtoDriftPackage +from drift_protocol.easy import DriftPackage, TriggerMessage, StatusCode + + +def test_build_package(): + """Test building package""" + message = TriggerMessage.build(1_000_000) + package = DriftPackage.build_from_msg( + message, + pkg_id=1_000_000, + status=StatusCode.GOOD, + source_timestamp=0, + publish_timestamp=2_000_000, + ) + + assert package.id > 0 + assert package.status == StatusCode.GOOD + assert package.source_timestamp == 0 + assert package.publish_timestamp == 2000.0 + + new_message = TriggerMessage() + package.unpack(new_message) + + assert new_message.timestamp == 1000.0 + + +def test_parse_package(): + """Should parse a protobuf message""" + proto = ProtoDriftPackage() + proto.id = 1_000_000 + + message = proto.SerializeToString() + package = DriftPackage.parse(message) + + assert package.id == 1_000_000 diff --git a/python/tests/easy/test_trigger.py b/python/tests/easy/test_trigger.py new file mode 100644 index 0000000..c963187 --- /dev/null +++ b/python/tests/easy/test_trigger.py @@ -0,0 +1,24 @@ +"""Unit tests for trigger.py""" + +from drift_protocol.easy import TriggerMessage + + +def test_build_trigger_message(): + """Test building trigger message""" + trigger_message = TriggerMessage.build(1_000_000) + assert trigger_message.timestamp == 1000.0 + + trigger_message = TriggerMessage.build(1000.0) + assert trigger_message.timestamp == 1000.0 + + trigger_message = TriggerMessage.build() + assert trigger_message.timestamp > 0 + + +def test_parse_trigger_message(): + """Test parsing trigger message""" + trigger_message = TriggerMessage.build(1_000_000) + message = trigger_message._proto.SerializeToString() + + parsed_trigger_message = TriggerMessage.parse(message) + assert parsed_trigger_message.timestamp == 1000.0 From 5dc7b9b407465cd978db9af3ef73fe6f5bca8feb Mon Sep 17 00:00:00 2001 From: Alexey Timin Date: Fri, 24 Nov 2023 16:05:33 +0100 Subject: [PATCH 2/2] add typed data support --- python/pkg/drift_protocol/easy/package.py | 72 +++++++++++++++++++---- python/pkg/drift_protocol/easy/trigger.py | 4 +- python/tests/easy/test_package.py | 40 +++++++++++++ 3 files changed, 104 insertions(+), 12 deletions(-) diff --git a/python/pkg/drift_protocol/easy/package.py b/python/pkg/drift_protocol/easy/package.py index f084a4b..4e2dc72 100644 --- a/python/pkg/drift_protocol/easy/package.py +++ b/python/pkg/drift_protocol/easy/package.py @@ -1,16 +1,19 @@ """Drfit Package""" import time -from typing import Optional, Dict +from typing import Optional, Dict, Tuple from google.protobuf.any_pb2 import Any -from drift_protocol.common import DriftPackage as ProtoDriftPackage, StatusCode, DataPayload -from drift_bytes import Variant, OutputBuffer +from drift_protocol.common import ( + DriftPackage as ProtoDriftPackage, + StatusCode, + DataPayload, +) +from drift_bytes import Variant, OutputBuffer, InputBuffer from drift_protocol.meta import TypedDataInfo, MetaInfo class DriftPackage: - def __init__(self): self._proto = ProtoDriftPackage() @@ -30,7 +33,12 @@ def parse(message): return package @staticmethod - def build_from_msg(message, pkg_id: Optional[int | float] = None, status: StatusCode = StatusCode.GOOD, **kwargs) -> "DriftPackage": + def build_from_msg( + message, + pkg_id: Optional[int | float] = None, + status: StatusCode = StatusCode.GOOD, + **kwargs + ) -> "DriftPackage": """ Build drift package @@ -77,10 +85,16 @@ def build_from_msg(message, pkg_id: Optional[int | float] = None, status: Status return package @staticmethod - def build_from_typed_data(values: Dict[str, Variant.SUPPORTED_TYPES], - statuses: Optional[Dict[str, StatusCode]] = None, - pkg_id: Optional[int | float] = None, status: StatusCode = StatusCode.GOOD, **kwargs) -> "DriftPackage": - + def build_typed_data( + values: Dict[str, Variant.SUPPORTED_TYPES], + statuses: Optional[Dict[str, int]] = None, + pkg_id: Optional[int | float] = None, + status: StatusCode = StatusCode.GOOD, + **kwargs + ) -> "DriftPackage": + """ + Build drift package from typed data + """ statuses = statuses or {} info = TypedDataInfo() @@ -92,16 +106,20 @@ def build_from_typed_data(values: Dict[str, Variant.SUPPORTED_TYPES], item = TypedDataInfo.Item() item.name = name - item.type = statuses.get(name, StatusCode.GOOD) + item.status = statuses.get(name, StatusCode.GOOD) info.items.append(item) data_payload = DataPayload() data_payload.data = buffer.bytes() - package = DriftPackage.build_from_msg(data_payload, pkg_id=pkg_id, status=status, **kwargs) + package = DriftPackage.build_from_msg( + data_payload, pkg_id=pkg_id, status=status, **kwargs + ) package._proto.meta.type = MetaInfo.TYPED_DATA package._proto.meta.typed_data_info.CopyFrom(info) + return package + @property def id(self) -> int: """ @@ -130,6 +148,38 @@ def source_timestamp(self) -> float: """ return self._proto.source_timestamp.ToMilliseconds() / 1000 + @property + def data_type(self): + """ + Data type + + Returns: + Data type or None if there is no meta info + """ + if self._proto.meta: + return self._proto.meta.type + return None + + def as_typed_data(self) -> Optional[Dict[str, Tuple[Variant.SUPPORTED_TYPES, int]]]: + """ + Get typed data + + Returns: + Typed data or None if there is no meta info + """ + if self._proto.meta and self._proto.meta.type == MetaInfo.TYPED_DATA: + data = {} + payload = DataPayload() + self._proto.data[0].Unpack(payload) + + buffer = InputBuffer(payload.data) + + for idx, item in enumerate(self._proto.meta.typed_data_info.items): + data[item.name] = (buffer[idx].value, item.status) + return data + + return None + def unpack(self, msg): """ Unpack message it must be either a protobuf message or an "easy" message diff --git a/python/pkg/drift_protocol/easy/trigger.py b/python/pkg/drift_protocol/easy/trigger.py index e3f12f1..fdc8067 100644 --- a/python/pkg/drift_protocol/easy/trigger.py +++ b/python/pkg/drift_protocol/easy/trigger.py @@ -2,7 +2,9 @@ import time from typing import Dict, Optional -from drift_protocol.trigger_service.trigger_message_pb2 import TriggerMessage as ProtoTriggerMessage +from drift_protocol.trigger_service.trigger_message_pb2 import ( + TriggerMessage as ProtoTriggerMessage, +) class TriggerMessage: diff --git a/python/tests/easy/test_package.py b/python/tests/easy/test_package.py index a0e1d74..e3ab5a0 100644 --- a/python/tests/easy/test_package.py +++ b/python/tests/easy/test_package.py @@ -2,6 +2,7 @@ from drift_protocol.common import DriftPackage as ProtoDriftPackage from drift_protocol.easy import DriftPackage, TriggerMessage, StatusCode +from drift_protocol.meta import MetaInfo def test_build_package(): @@ -35,3 +36,42 @@ def test_parse_package(): package = DriftPackage.parse(message) assert package.id == 1_000_000 + + +def test_build_from_typed_data(): + """Should build a package from typed data""" + values = { + "int": 100, + "float": 25.9, + "string": "test", + "bool": True, + "list": [1, 2, 3], + } + + statuses = { + "string": StatusCode.BAD, + "bool": StatusCode.GOOD, + } + + package = DriftPackage.build_typed_data( + values, + statuses, + pkg_id=1_000_000, + status=StatusCode.GOOD, + source_timestamp=0, + publish_timestamp=2_000_000, + ) + + assert package.id > 0 + assert package.status == StatusCode.GOOD + assert package.source_timestamp == 0 + assert package.publish_timestamp == 2000.0 + + assert package.data_type == MetaInfo.TYPED_DATA + assert package.as_typed_data() == { + "int": (100, StatusCode.GOOD), + "float": (25.9, StatusCode.GOOD), + "string": ("test", StatusCode.BAD), + "bool": (True, StatusCode.GOOD), + "list": ([1, 2, 3], StatusCode.GOOD), + }