Source code for go2.modules.input.input_module

from typing import Callable, Optional
from typing_extensions import override

from .callback_manager import InputSignalCallbackManager, UnitreeRemoteControllerInputParser
from .controller_state import ControllerState
from .input_signal import InputSignal

from ...core.module import DogModule
from ...communication.dds import DDSTopics



class InputModule(DogModule):
    """
    Users should not access or construct this class directly.
    Rather, they should access it through the :class:`~core.controller.Go2Controller` instance.
    """
    def __init__(self) -> None:
        super().__init__("Input")

    @override
    def _initialize(self) -> None:
        """
        Set up input parsing and callback management. This is called internally,
        and should not be called directly by users.

        Notes
        -----
        - If `use_sdk` is False, no live input is initialized
        - Internal: subscribes to DDS LOW_STATE topic for live controller messages
        """
        from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowState_
        from unitree_sdk2py.core.channel import ChannelSubscriber

        self._input_parser = UnitreeRemoteControllerInputParser()
        self._callback_manager = InputSignalCallbackManager()

        self._lowstate_subscriber = ChannelSubscriber(DDSTopics.LOW_STATE, LowState_)
        self._lowstate_subscriber.Init(self._process_input, 10)



    def register_callback(
        self,
        signal: InputSignal,
        callback: Callable[[ControllerState], None],
        name: Optional[str] = None,
        threshold: float = 0.1
    ):
        """
        Register a callback for a specific controller signal.

        Parameters
        ----------
        signal : InputSignal
            The input signal to monitor (e.g., BUTTON_A, LEFT_STICK_X)
        callback : Callable[[ControllerState], None]
            Function to call when the signal triggers
        name : str, optional
            Optional human-readable name for the callback
        threshold : float, optional
            Threshold for analog inputs (default 0.1)

        Returns
        -------
        Callback
            Registered callback object
        """
        return self._callback_manager._register(signal, callback, name, threshold)

    


    def unregister_callback(
        self,
        signal: InputSignal,
        callback: Callable[[ControllerState], None]
    ) -> None:
        """
        Unregister a previously registered callback.

        Parameters
        ----------
        signal : InputSignal
            The signal whose callback should be removed
        callback : Callable[[ControllerState], None]
            The previously registered callback function
        """
        self._callback_manager._unregister(signal, callback)

    
    @override
    def _shutdown(self) -> None:
        """
        Clean up input resources. This is handled automatically and shouldn't be called by users.

        Notes
        -----
        - Stops DDS subscription and clears all callbacks
        - Should be called when input is no longer needed
        """
        if self._lowstate_subscriber:
            self._lowstate_subscriber.Close()
        if self._callback_manager:
            self._callback_manager._shutdown()

    def _process_input(self, msg) -> ControllerState:
        """
        Process incoming controller messages.

        Parameters
        ----------
        msg : object
            DDS message containing raw controller data

        Returns
        -------
        ControllerState
            Current controller state after parsing
        """
        controller_state = self._input_parser._parse(msg.wireless_remote)
        self._callback_manager._handle(controller_state)

        return controller_state