behavior-tree/tests/helpers/ros_shim.py

"""Minimal ROS runtime and message classes generated from repository IDL."""
from __future__ import annotations

import importlib.util
import pathlib
import sys
import types

ROOT = pathlib.Path(__file__).resolve().parents[2]
IDL = ROOT / "ros2" / "bt_skill_interfaces"


def _module(name):
    module = types.ModuleType(name)
    sys.modules[name] = module
    return module


def _fields(path, section=0):
    parts = [[]]
    for raw in path.read_text().splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        if line == "---":
            parts.append([])
        else:
            parts[-1].append(line)
    return parts[section]


class RosClock:
    _nanoseconds = 10_000_000_000

    def now(self):
        self._nanoseconds += 1_000_000
        return types.SimpleNamespace(
            nanoseconds=self._nanoseconds,
            to_msg=lambda: _type("builtin_interfaces/Time")(
                sec=self._nanoseconds // 1_000_000_000,
                nanosec=self._nanoseconds % 1_000_000_000,
            ),
        )


_registry = {}


def _type(type_name):
    if type_name.endswith("[]"):
        return list
    return _registry[type_name]


def _default(type_name):
    if type_name.endswith("[]"):
        return []
    if type_name in ("string",):
        return ""
    if type_name in ("bool",):
        return False
    if type_name.startswith(("uint", "int", "float")):
        return 0
    return _type(type_name)()


def _make_class(name, lines):
    constants, fields = {}, []
    for line in lines:
        type_name, declaration = line.split(None, 1)
        if "=" in declaration:
            key, value = declaration.split("=", 1)
            constants[key] = int(value)
        else:
            fields.append((type_name, declaration))
    slots = tuple(field for _, field in fields)

    def init(self, **kwargs):
        for type_name, field in fields:
            setattr(self, field, kwargs.pop(field, _default(type_name)))
        if kwargs:
            raise TypeError("unexpected fields: " + ", ".join(kwargs))

    attrs = {"__slots__": slots, "__init__": init, **constants,
             "get_fields_and_field_types": classmethod(lambda cls: {field: kind for kind, field in fields})}
    return type(name, (), attrs)


def install():
    """Install deterministic rclpy and generated interface modules."""
    _registry.clear()
    for name in list(sys.modules):
        if name == "rclpy" or name.startswith("rclpy.") or name.startswith(("bt_skill_interfaces", "navigation_interfaces")):
            del sys.modules[name]

    builtin = _module("builtin_interfaces")
    builtin_msg = _module("builtin_interfaces.msg")
    builtin.msg = builtin_msg
    for name, fields in {
        "Time": ["int32 sec", "uint32 nanosec"],
        "Duration": ["int32 sec", "uint32 nanosec"],
    }.items():
        cls = _make_class(name, fields)
        setattr(builtin_msg, name, cls)
        _registry[f"builtin_interfaces/{name}"] = cls

    std = _module("std_msgs")
    std_msg = _module("std_msgs.msg")
    std.msg = std_msg
    header = _make_class("Header", ["builtin_interfaces/Time stamp", "string frame_id"])
    std_msg.Header = header
    _registry["std_msgs/Header"] = header
    string = _make_class("String", ["string data"])
    std_msg.String = string
    _registry["std_msgs/String"] = string

    geometry = _module("geometry_msgs")
    geometry_msg = _module("geometry_msgs.msg")
    geometry.msg = geometry_msg
    definitions = {
        "Point": ["float64 x", "float64 y", "float64 z"],
        "Quaternion": ["float64 x", "float64 y", "float64 z", "float64 w"],
        "Pose": ["geometry_msgs/Point position", "geometry_msgs/Quaternion orientation"],
        "PoseStamped": ["std_msgs/Header header", "geometry_msgs/Pose pose"],
        "PointStamped": ["std_msgs/Header header", "geometry_msgs/Point point"],
    }
    for name, fields in definitions.items():
        cls = _make_class(name, fields)
        setattr(geometry_msg, name, cls)
        _registry[f"geometry_msgs/{name}"] = cls

    for package_name in ("bt_skill_interfaces", "navigation_interfaces"):
        package_idl = IDL.parent / package_name
        package = _module(package_name)
        msg_module = _module(f"{package_name}.msg")
        package.msg = msg_module
        pending = {p.stem: _fields(p) for p in (package_idl / "msg").glob("*.msg")}
        while pending:
            progress = False
            for name, lines in list(pending.items()):
                deps = [line.split()[0].removesuffix("[]") for line in lines if "=" not in line]
                if all(dep in _registry or dep in ("string", "bool") or dep.startswith(("uint", "int", "float")) for dep in deps):
                    cls = _make_class(name, lines)
                    setattr(msg_module, name, cls)
                    _registry[f"{package_name}/{name}"] = cls
                    del pending[name]
                    progress = True
            if not progress:
                raise RuntimeError("unresolved IDL: " + repr(pending))

        action_module = _module(f"{package_name}.action")
        package.action = action_module
        for path in (package_idl / "action").glob("*.action"):
            action = type(path.stem, (), {})
            action.Goal = _make_class("Goal", _fields(path, 0))
            action.Result = _make_class("Result", _fields(path, 1))
            action.Feedback = _make_class("Feedback", _fields(path, 2))
            setattr(action_module, path.stem, action)

        srv_module = _module(f"{package_name}.srv")
        package.srv = srv_module
        for path in (package_idl / "srv").glob("*.srv"):
            srv = type(path.stem, (), {})
            srv.Request = _make_class("Request", _fields(path, 0))
            srv.Response = _make_class("Response", _fields(path, 1))
            setattr(srv_module, path.stem, srv)

    rclpy = _module("rclpy")
    rclpy.ok = lambda: True
    rclpy.init = lambda **_: None
    rclpy.shutdown = lambda: None
    action = _module("rclpy.action")
    action.GoalResponse = types.SimpleNamespace(ACCEPT=1, REJECT=2)
    action.CancelResponse = types.SimpleNamespace(ACCEPT=1)
    action.ActionServer = lambda *args, **kwargs: types.SimpleNamespace(destroy=lambda: None)
    callbacks = _module("rclpy.callback_groups")
    callbacks.ReentrantCallbackGroup = object
    executors = _module("rclpy.executors")
    executors.MultiThreadedExecutor = lambda **_: types.SimpleNamespace(
        add_node=lambda node: None, spin=lambda: None, shutdown=lambda **_: None)
    node_module = _module("rclpy.node")

    class Node:
        parameters = {}

        def __init__(self, *_args, namespace="/sim/robot_01", **_kwargs):
            self._namespace = namespace
            self._clock = RosClock()

        def get_namespace(self): return self._namespace
        def declare_parameter(self, name, default): self.parameters.setdefault(name, default)
        def get_parameter(self, name): return types.SimpleNamespace(value=self.parameters[name])
        def resolve_topic_name(self, name): return self._namespace + "/" + name
        def create_publisher(self, _type, name, *_a, **_k):
            publisher = types.SimpleNamespace(name=name, values=[])
            publisher.publish = publisher.values.append
            return publisher
        def create_service(self, _type, name, callback, **_k): return types.SimpleNamespace(name=name, callback=callback)
        def create_timer(self, *_a, **_k): return object()
        def get_logger(self): return types.SimpleNamespace(warning=lambda *_: None, error=lambda *_: None)
        def get_clock(self): return self._clock
        def destroy_node(self): pass

    node_module.Node = Node
    rclpy.action, rclpy.callback_groups, rclpy.executors, rclpy.node = action, callbacks, executors, node_module


def load_mock_module(parameters=None):
    install()
    sys.modules["rclpy.node"].Node.parameters = dict(parameters or {})
    package = _module("bt_mock_servers")
    package.__path__ = [str(IDL.parent / "bt_mock_servers" / "bt_mock_servers")]
    source = IDL.parent / "bt_mock_servers" / "bt_mock_servers" / "mock_skills.py"
    spec = importlib.util.spec_from_file_location("bt_mock_servers.mock_skills", source)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class GoalHandle:
    def __init__(self, request, cancel=False, uuid=bytes(range(16))):
        self.request, self.is_cancel_requested = request, cancel
        self.goal_id = types.SimpleNamespace(uuid=uuid)
        self.feedback, self.native = [], None
        self.is_active = True
        self.cancel_acknowledged = False

    def execute(self): pass
    def publish_feedback(self, value): self.feedback.append(value)
    def succeed(self): self.native, self.is_active = "succeeded", False
    def abort(self): self.native, self.is_active = "aborted", False
    def canceled(self): self.native, self.is_active = "canceled", False

    def request_cancel(self):
        self.is_cancel_requested = True
        self.cancel_acknowledged = True
        return True