ros_builder.py

"""Builders for auto-generated ROS1 and ROS2 package wrappers."""

import opengen.definitions as og_dfn

import datetime
import logging
import os
import shutil
import sys

import jinja2


def make_dir_if_not_exists(directory):
    """Create ``directory`` if it does not already exist.

    :param directory: Path to the directory to create.
    :type directory: str
    """
    if not os.path.exists(directory):
        os.makedirs(directory)


def get_ros_template(template_subdir, name):
    """Load a Jinja template from a ROS-specific template subdirectory.

    :param template_subdir: Template subdirectory name, e.g. ``"ros"`` or
        ``"ros2"``.
    :type template_subdir: str
    :param name: Template file name.
    :type name: str

    :return: Loaded Jinja template.
    :rtype: jinja2.Template
    """
    file_loader = jinja2.FileSystemLoader(og_dfn.templates_subdir(template_subdir))
    env = jinja2.Environment(loader=file_loader, autoescape=True)
    return env.get_template(name)


class _BaseRosBuilder:
    """
    Shared code generation logic for ROS-related packages.

    This base class contains the common file-generation pipeline used by both
    :class:`RosBuilder` and :class:`ROS2Builder`. Subclasses specialize the
    process by providing the package configuration object, template
    subdirectory, launch file name, and final user-facing instructions.

    :ivar _meta: Optimizer metadata used to render the package templates.
    :ivar _build_config: Global build configuration for the generated solver.
    :ivar _solver_config: Solver configuration used when rendering node code.
    :ivar _logger: Logger dedicated to the concrete builder implementation.
    """

    #: Template subdirectory under ``opengen/templates`` used by the builder.
    _template_subdir = None
    #: Fully-qualified logger name for the concrete builder.
    _logger_name = None
    #: Short logger tag shown in log messages.
    _logger_tag = None
    #: Launch file generated by the concrete builder.
    _launch_file_name = None

    def __init__(self, meta, build_config, solver_config):
        """Initialise a shared ROS package builder.

        :param meta: Optimizer metadata.
        :param build_config: Build configuration object.
        :param solver_config: Solver configuration object.
        """
        self._meta = meta
        self._build_config = build_config
        self._solver_config = solver_config
        self._logger = logging.getLogger(self._logger_name)
        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(1)
        c_format = logging.Formatter(
            f'[%(levelname)s] <<{self._logger_tag}>> %(message)s')
        stream_handler.setFormatter(c_format)
        self._logger.setLevel(1)
        self._logger.handlers.clear()
        self._logger.addHandler(stream_handler)
        self._logger.propagate = False

    @property
    def _ros_config(self):
        """Return the ROS/ROS2 package configuration for the subclass.

        :return: ROS configuration object used by the concrete builder.
        :raises NotImplementedError: If a subclass does not provide this hook.
        """
        raise NotImplementedError

    def _template(self, name):
        """Return a template from the builder's template subdirectory.

        :param name: Template file name.
        :type name: str

        :return: Loaded Jinja template.
        :rtype: jinja2.Template
        """
        return get_ros_template(self._template_subdir, name)

    def _target_dir(self):
        """Return the root directory of the generated optimizer project.

        :return: Absolute path to the generated optimizer directory.
        :rtype: str
        """
        return os.path.abspath(
            os.path.join(
                self._build_config.build_dir,
                self._meta.optimizer_name))

    def _ros_target_dir(self):
        """Return the root directory of the generated ROS package.

        :return: Absolute path to the generated ROS/ROS2 package directory.
        :rtype: str
        """
        return os.path.abspath(
            os.path.join(
                self._build_config.build_dir,
                self._meta.optimizer_name,
                self._ros_config.package_name))

    def _generate_ros_dir_structure(self):
        """Create the directory structure for the generated ROS package."""
        self._logger.info("Generating directory structure")
        target_ros_dir = self._ros_target_dir()
        make_dir_if_not_exists(target_ros_dir)
        for directory_name in ('include', 'extern_lib', 'src', 'msg', 'config', 'launch'):
            make_dir_if_not_exists(os.path.abspath(
                os.path.join(target_ros_dir, directory_name)))

    def _generate_ros_package_xml(self):
        """Render and write ``package.xml`` for the generated package."""
        self._logger.info("Generating package.xml")
        target_ros_dir = self._ros_target_dir()
        template = self._template('package.xml')
        output_template = template.render(meta=self._meta, ros=self._ros_config)
        target_rospkg_path = os.path.join(target_ros_dir, "package.xml")
        with open(target_rospkg_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _generate_ros_cmakelists(self):
        """Render and write the package ``CMakeLists.txt`` file."""
        self._logger.info("Generating CMakeLists")
        target_ros_dir = self._ros_target_dir()
        template = self._template('CMakeLists.txt')
        output_template = template.render(meta=self._meta, ros=self._ros_config)
        target_rospkg_path = os.path.join(target_ros_dir, "CMakeLists.txt")
        with open(target_rospkg_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _copy_ros_files(self):
        """Copy generated bindings, static library, and message files."""
        self._logger.info("Copying external dependencies")
        target_ros_dir = self._ros_target_dir()

        header_file_name = self._meta.optimizer_name + '_bindings.hpp'
        target_include_filename = os.path.abspath(
            os.path.join(target_ros_dir, 'include', header_file_name))
        original_include_file = os.path.abspath(
            os.path.join(self._target_dir(), header_file_name))
        shutil.copyfile(original_include_file, target_include_filename)

        if sys.platform == "win32":
            lib_file_name = self._meta.optimizer_name + '.lib'
        else:
            lib_file_name = 'lib' + self._meta.optimizer_name + '.a'
        target_lib_file_name = os.path.abspath(
            os.path.join(target_ros_dir, 'extern_lib', lib_file_name))
        original_lib_file = os.path.abspath(
            os.path.join(
                self._target_dir(),
                'target',
                self._build_config.build_mode,
                lib_file_name))
        shutil.copyfile(original_lib_file, target_lib_file_name)

        for message_name in ('OptimizationParameters.msg', 'OptimizationResult.msg'):
            original_message = os.path.abspath(
                os.path.join(
                    og_dfn.templates_dir(),
                    self._template_subdir,
                    message_name))
            target_message = os.path.abspath(
                os.path.join(target_ros_dir, 'msg', message_name))
            shutil.copyfile(original_message, target_message)

    def _generate_ros_params_file(self):
        """Render and write the runtime parameter YAML file."""
        self._logger.info("Generating open_params.yaml")
        target_ros_dir = self._ros_target_dir()
        template = self._template('open_params.yaml')
        output_template = template.render(meta=self._meta, ros=self._ros_config)
        target_yaml_fname = os.path.join(target_ros_dir, "config", "open_params.yaml")
        with open(target_yaml_fname, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _generate_ros_node_header(self):
        """Render and write the generated node header file."""
        self._logger.info("Generating open_optimizer.hpp")
        target_ros_dir = self._ros_target_dir()
        template = self._template('open_optimizer.hpp')
        output_template = template.render(
            meta=self._meta,
            ros=self._ros_config,
            solver_config=self._solver_config)
        target_rosnode_header_path = os.path.join(
            target_ros_dir, "include", "open_optimizer.hpp")
        with open(target_rosnode_header_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _generate_ros_node_cpp(self):
        """Render and write the generated node implementation file."""
        self._logger.info("Generating open_optimizer.cpp")
        target_ros_dir = self._ros_target_dir()
        template = self._template('open_optimizer.cpp')
        output_template = template.render(
            meta=self._meta,
            ros=self._ros_config,
            timestamp_created=datetime.datetime.now())
        target_rosnode_cpp_path = os.path.join(target_ros_dir, "src", "open_optimizer.cpp")
        with open(target_rosnode_cpp_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _generate_ros_launch_file(self):
        """Render and write the package launch file."""
        self._logger.info("Generating %s", self._launch_file_name)
        target_ros_dir = self._ros_target_dir()
        template = self._template(self._launch_file_name)
        output_template = template.render(meta=self._meta, ros=self._ros_config)
        target_rosnode_launch_path = os.path.join(
            target_ros_dir, "launch", self._launch_file_name)
        with open(target_rosnode_launch_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _generate_ros_readme_file(self):
        """Render and write the generated package README."""
        self._logger.info("Generating README.md")
        target_ros_dir = self._ros_target_dir()
        template = self._template('README.md')
        output_template = template.render(ros=self._ros_config)
        target_readme_path = os.path.join(target_ros_dir, "README.md")
        with open(target_readme_path, "w", encoding="utf-8") as fh:
            fh.write(output_template)

    def _symbolic_link_info_message(self):
        """Emit final user-facing setup instructions for the generated package.

        :raises NotImplementedError: If a subclass does not provide this hook.
        """
        raise NotImplementedError

    def build(self):
        """
        Generate all ROS/ROS2 wrapper files for the current optimizer.

        This method creates the package directory structure, copies the
        generated solver artefacts, renders all templates, and logs final setup
        instructions for the user.
        """
        self._generate_ros_dir_structure()
        self._generate_ros_package_xml()
        self._generate_ros_cmakelists()
        self._copy_ros_files()
        self._generate_ros_params_file()
        self._generate_ros_node_header()
        self._generate_ros_node_cpp()
        self._generate_ros_launch_file()
        self._generate_ros_readme_file()
        self._symbolic_link_info_message()


class RosBuilder(_BaseRosBuilder):
    """
    Builder for ROS1 package generation.

    This specialization uses the ``templates/ros`` template set and the
    ROS1-specific configuration stored in
    :attr:`opengen.config.build_config.BuildConfiguration.ros_config`.
    """

    _template_subdir = 'ros'
    _logger_name = 'opengen.builder.RosBuilder'
    _logger_tag = 'ROS'
    _launch_file_name = 'open_optimizer.launch'

    @property
    def _ros_config(self):
        """Return the ROS1 package configuration."""
        return self._build_config.ros_config

    def _symbolic_link_info_message(self):
        """Log the final ROS1 workspace integration instructions."""
        target_ros_dir = self._ros_target_dir()
        self._logger.info("ROS package was built successfully. Now run:")
        self._logger.info("ln -s %s  ~/catkin_ws/src/", target_ros_dir)
        self._logger.info("cd ~/catkin_ws/; catkin_make")


class ROS2Builder(_BaseRosBuilder):
    """
    Builder for ROS2 package generation.

    This specialization uses the ``templates/ros2`` template set and the
    ROS2-specific configuration stored in
    :attr:`opengen.config.build_config.BuildConfiguration.ros2_config`.
    """

    _template_subdir = 'ros2'
    _logger_name = 'opengen.builder.ROS2Builder'
    _logger_tag = 'ROS2'
    _launch_file_name = 'open_optimizer.launch.py'

    @property
    def _ros_config(self):
        """Return the ROS2 package configuration."""
        return self._build_config.ros2_config

    def _symbolic_link_info_message(self):
        """Log the final ROS2 workspace integration instructions."""
        target_ros_dir = self._ros_target_dir()
        self._logger.info("ROS2 package was built successfully. Now run:")
        self._logger.info("ln -s %s  ~/ros2_ws/src/", target_ros_dir)
        self._logger.info("cd ~/ros2_ws/; colcon build --packages-select %s",
                          self._ros_config.package_name)