plexus-python-common 1.0.52__tar.gz → 1.0.53__tar.gz

{plexus_python_common-1.0.52 → plexus_python_common-1.0.53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plexus-python-common
-Version: 1.0.52
+Version: 1.0.53
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13

plexus_python_common-1.0.53/src/plexus/common/utils/bagutils.py

@@ -0,0 +1,329 @@
+import contextlib
+import os
+from collections.abc import Generator
+from contextlib import AbstractContextManager
+from typing import Literal
+
+import sqlalchemy as sa
+import sqlalchemy.orm as sa_orm
+from iker.common.utils.pathutils import make_path
+
+__all__ = [
+    "SerializationFormat",
+    "BagSchema",
+    "BagMetadata",
+    "BagMessageDefinition",
+    "BagTopic",
+    "BagMessage",
+    "BagReader",
+    "BagWriter",
+    "bag_reader",
+    "bag_writer",
+]
+
+SerializationFormat = Literal["cdr", "cdr2", "json", "yaml"]
+
+default_bag_db_file = "bag.db"
+
+
+class BagBase(sa_orm.DeclarativeBase):
+    pass
+
+
+class BagSchema(BagBase):
+    __tablename__ = "schema"
+
+    schema_version: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True)
+    ros_distro: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+
+
+class BagMetadata(BagBase):
+    __tablename__ = "metadata"
+
+    id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True, autoincrement=True)
+    metadata_version: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, nullable=False, unique=True)
+    metadata_text: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+
+
+class BagMessageDefinition(BagBase):
+    __tablename__ = "message_definitions"
+
+    id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True, autoincrement=True)
+    message_type: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False, unique=True)
+    encoding: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+    encoded_message_definition: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+    type_description_hash: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+
+    topics: sa_orm.Mapped[list["BagTopic"]] = sa_orm.relationship("BagTopic", back_populates="message_definition")
+
+
+class BagTopic(BagBase):
+    __tablename__ = "topics"
+
+    id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True, autoincrement=True)
+    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False, unique=True)
+    message_definition_id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer,
+                                                                     sa.ForeignKey("message_definitions.id"),
+                                                                     nullable=False)
+    serialization_format: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+    type_description_hash: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.String, nullable=False)
+
+    messages: sa_orm.Mapped[list["BagMessage"]] = sa_orm.relationship("BagMessage", back_populates="topic")
+    message_definition: sa_orm.Mapped["BagMessageDefinition"] = sa_orm.relationship("BagMessageDefinition",
+                                                                                    back_populates="topics")
+
+
+class BagMessage(BagBase):
+    __tablename__ = "messages"
+
+    id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True, autoincrement=True)
+    topic_id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, sa.ForeignKey("topics.id"), nullable=False)
+    timestamp: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, nullable=False)
+    data: sa_orm.Mapped[bytes] = sa_orm.mapped_column(sa.LargeBinary, nullable=False)
+
+    topic: sa_orm.Mapped["BagTopic"] = sa_orm.relationship("BagTopic", back_populates="messages")
+
+
+class BagIOBase(object):
+    def __init__(self, bag_file_path: str | os.PathLike[str]):
+        self.bag_file_path = make_path(bag_file_path)
+
+        self.engine: sa.Engine | None = None
+        self.session: sa_orm.Session | None = None
+
+        self.message_definitions_lookup: dict[str, int] = {}
+        self.topics_lookup: dict[str, int] = {}
+
+    @contextlib.contextmanager
+    def setup(self):
+        self.engine = sa.create_engine(f"sqlite:///{str(self.bag_file_path.absolute())}")
+        self.session = sa_orm.Session(self.engine)
+        try:
+            yield
+        finally:
+            self.message_definitions_lookup = {db_message_definition.message_type: db_message_definition.id
+                                               for db_message_definition in self.message_definitions()}
+            self.topics_lookup = {db_topic.name: db_topic.id for db_topic in self.topics()}
+
+    @contextlib.contextmanager
+    def teardown(self):
+        try:
+            yield
+        finally:
+            self.session.close()
+            self.engine.dispose()
+
+    def get_metadata(self, metadata_version: int) -> BagMetadata | None:
+        return (
+            self
+            .session
+            .query(BagMetadata)
+            .filter(BagMetadata.metadata_version == metadata_version)
+            .one_or_none()
+        )
+
+    def get_message_definition(self, message_type: str) -> BagMessageDefinition | None:
+        return (
+            self
+            .session
+            .query(BagMessageDefinition)
+            .filter(BagMessageDefinition.message_type == message_type)
+            .one_or_none()
+        )
+
+    def get_topic(self, name: str) -> BagTopic | None:
+        return self.session.query(BagTopic).filter(BagTopic.name == name).one_or_none()
+
+    def metadata(self) -> list[BagMetadata]:
+        return self.session.query(BagMetadata).all()
+
+    def message_definitions(self) -> list[BagMessageDefinition]:
+        return self.session.query(BagMessageDefinition).all()
+
+    def topics(self) -> list[BagTopic]:
+        return self.session.query(BagTopic).all()
+
+
+class BagReader(AbstractContextManager, BagIOBase):
+    @staticmethod
+    def open_ros2(db_dir: str | os.PathLike[str], db_filename: str = default_bag_db_file, **kwargs) -> "BagReader":
+        """
+        Creates a BagReader instance to read messages from a ROS2 bag file in the specified directory.
+
+        :param db_dir: path to the directory where the SQLite DB file is located.
+        :param db_filename: name of the SQLite DB file to read. Default is "bag.db".
+        :param kwargs: additional keyword arguments to pass to the BagReader constructor.
+        :return: a BagReader instance.
+        """
+        return BagReader(make_path(db_dir) / db_filename)
+
+    def __init__(self, bag_file_path: str | os.PathLike[str]):
+        """
+        Creates a BagReader instance to read messages from a ROS2 bag file.
+
+        :param bag_file_path: path to the SQLite DB file to read.
+        """
+        super().__init__(bag_file_path)
+
+        if not self.bag_file_path.exists():
+            raise FileNotFoundError(f"could not find SQLite DB at '{str(self.bag_file_path)}'")
+
+    def __enter__(self):
+        with self.setup():
+            pass
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        with self.teardown():
+            pass
+
+    def iter_messages(
+        self,
+        topic_names: list[str] | None = None,
+        begin_timestamp: int | None = None,
+        end_timestamp: int | None = None,
+        *,
+        batch_size: int = 1000,
+    ) -> Generator[BagMessage, None, None]:
+        query = self.session.query(BagMessage)
+        if topic_names is not None:
+            topic_ids = [topic_id for topic_name, topic_id in self.topics_lookup.items() if topic_name in topic_names]
+            query = query.where(BagMessage.topic_id.in_(topic_ids))
+        if begin_timestamp is not None:
+            query = query.where(BagMessage.timestamp >= begin_timestamp)
+        if end_timestamp is not None:
+            query = query.where(BagMessage.timestamp <= end_timestamp)
+
+        for db_message in query.order_by(BagMessage.timestamp.asc()).yield_per(batch_size):
+            yield db_message
+
+
+class BagWriter(AbstractContextManager, BagIOBase):
+    @staticmethod
+    def open_ros2(db_dir: str | os.PathLike[str], db_filename: str = default_bag_db_file, **kwargs) -> "BagWriter":
+        """
+        Creates a BagWriter instance to write messages to a ROS2 bag file in the specified directory.
+
+        :param db_dir: path to the directory where the SQLite DB file will be created.
+        :param db_filename: name of the SQLite DB file to create. Default is "bag.db".
+        :param kwargs: additional keyword arguments to pass to the BagWriter constructor.
+        :return: a BagWriter instance.
+        """
+        return BagWriter(make_path(db_dir) / db_filename, **kwargs)
+
+    def __init__(self, bag_file_path: str | os.PathLike[str], *, overwrite: bool = True, exist_ok: bool = False):
+        """
+        Creates a BagWriter instance to write messages to a ROS2 bag file.
+
+        :param bag_file_path: path to the SQLite DB file to create.
+        :param overwrite: whether to overwrite the SQLite DB file if it already exists. If False and the file already
+        exists, a FileExistsError is raised. If True, the existing file is deleted and a new one is created.
+        :param exist_ok: whether to ignore if the SQLite DB file already exists. If False and the file already exists,
+        a FileExistsError is raised. If True, the existing file is used and no new file is created.
+        """
+        super().__init__(bag_file_path)
+
+        if not self.bag_file_path.parent.exists():
+            self.bag_file_path.parent.mkdir(parents=True, exist_ok=True)
+        if self.bag_file_path.exists():
+            if not overwrite and not exist_ok:
+                raise FileExistsError(f"SQLite DB file already exists at '{str(self.bag_file_path)}'")
+            if overwrite:
+                self.bag_file_path.unlink()
+
+    def __enter__(self):
+        with self.setup():
+            BagBase.metadata.create_all(self.engine)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        with self.teardown():
+            if exc_type is None:
+                self.session.commit()
+            else:
+                self.session.rollback()
+
+    def add_metadata(self, metadata_version: int, metadata_text: str, *, return_existing: bool = False) -> BagMetadata:
+        if (db_metadata := self.get_metadata(metadata_version)) is not None:
+            if return_existing:
+                return db_metadata
+            raise ValueError(f"metadata with metadata_version '{metadata_version}' already exists in the bag file")
+
+        db_metadata = BagMetadata(metadata_version=metadata_version, metadata_text=metadata_text)
+        self.session.add(db_metadata)
+        self.session.flush()
+
+        return db_metadata
+
+    def add_message_definition(
+        self,
+        message_type: str,
+        encoding: str = "",
+        encoded_message_definition: str = "",
+        type_description_hash: str = "",
+        *,
+        return_existing: bool = False,
+    ) -> BagMessageDefinition:
+        if (db_message_definition := self.get_message_definition(message_type)) is not None:
+            if return_existing:
+                return db_message_definition
+            raise ValueError(f"message definition with message_type '{message_type}' already exists in the bag file")
+
+        db_message_definition = BagMessageDefinition(
+            message_type=message_type,
+            encoding=encoding,
+            encoded_message_definition=encoded_message_definition,
+            type_description_hash=type_description_hash,
+        )
+        self.session.add(db_message_definition)
+        self.session.flush()
+
+        self.message_definitions_lookup[db_message_definition.message_type] = db_message_definition.id
+
+        return db_message_definition
+
+    def add_topic(
+        self,
+        name: str,
+        message_type: str,
+        serialization_format: SerializationFormat = "cdr",
+        type_description_hash: str = "",
+        *,
+        return_existing: bool = False,
+    ) -> BagTopic:
+        message_definition_id = self.message_definitions_lookup.get(message_type)
+        if message_definition_id is None:
+            raise ValueError(f"message definition with message_type '{message_type}' does not exist in the bag file")
+
+        if (db_topic := self.get_topic(name)) is not None:
+            if return_existing:
+                return db_topic
+            raise ValueError(f"topic with name '{name}' already exists in the bag file")
+
+        db_topic = BagTopic(
+            name=name,
+            message_definition_id=message_definition_id,
+            serialization_format=serialization_format,
+            type_description_hash=type_description_hash,
+        )
+        self.session.add(db_topic)
+        self.session.flush()
+
+        self.topics_lookup[db_topic.name] = db_topic.id
+
+        return db_topic
+
+    def write_message(self, name: str, timestamp: int, data: bytes) -> BagMessage:
+        topic_id = self.topics_lookup.get(name)
+        if topic_id is None:
+            raise ValueError(f"topic with name '{name}' does not exist in the bag file")
+
+        db_message = BagMessage(topic_id=topic_id, timestamp=timestamp, data=data)
+        self.session.add(db_message)
+        self.session.flush()
+
+        return db_message
+
+
+bag_reader = BagReader
+bag_writer = BagWriter