RSDK-807: Extra params for Input controller (#180)

Co-authored-by: viambot <[email protected]>

Author: cheukt

examples/server/v1/components.py

@@ -395,7 +395,7 @@ def _execute_callback(self, event):
             if callback:
                 callback(event)
 
-    async def get_controls(self, **kwargs) -> List[Control]:
+    async def get_controls(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Control]:
         return [
             Control.ABSOLUTE_X,
             Control.ABSOLUTE_Y,
@@ -418,7 +418,7 @@ async def get_controls(self, **kwargs) -> List[Control]:
             Control.BUTTON_MENU,
         ]
 
-    async def get_events(self, **kwargs) -> Dict[Control, Event]:
+    async def get_events(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Dict[Control, Event]:
         with self.lock:
             return {key: value for (key, value) in self.last_events.items()}
 

src/viam/components/input/client.py

@@ -3,6 +3,7 @@
 from time import time
 from typing import Any, Dict, List, Optional
 from grpclib import GRPCError, Status
+from google.protobuf.struct_pb2 import Struct
 
 from grpclib.client import Channel
 import viam
@@ -19,6 +20,7 @@
     StreamEventsResponse,
     TriggerEventRequest,
 )
+from viam.utils import dict_to_struct
 
 from .input import Control, ControlFunction, Controller, Event, EventType
 
@@ -36,19 +38,29 @@ def __init__(self, name: str, channel: Channel):
         self._stream_lock = Lock()
         self._is_streaming = False
         self._is_stream_ready = False
+        self._callback_extra: Struct = dict_to_struct({})
         super().__init__(name)
 
-    async def get_controls(self, *, timeout: Optional[float] = None) -> List[Control]:
-        request = GetControlsRequest(controller=self.name)
+    async def get_controls(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Control]:
+        if extra is None:
+            extra = {}
+        request = GetControlsRequest(controller=self.name, extra=dict_to_struct(extra))
         response: GetControlsResponse = await self.client.GetControls(request, timeout=timeout)
         return [Control(control) for control in response.controls]
 
-    async def get_events(self, *, timeout: Optional[float] = None) -> Dict[Control, Event]:
-        request = GetEventsRequest(controller=self.name)
+    async def get_events(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Dict[Control, Event]:
+        if extra is None:
+            extra = {}
+        request = GetEventsRequest(controller=self.name, extra=dict_to_struct(extra))
         response: GetEventsResponse = await self.client.GetEvents(request, timeout=timeout)
         return {Control(event.control): Event.from_proto(event) for (event) in response.events}
 
-    def register_control_callback(self, control: Control, triggers: List[EventType], function: Optional[ControlFunction]):
+    def register_control_callback(
+        self, control: Control, triggers: List[EventType], function: Optional[ControlFunction], extra: Optional[Dict[str, Any]] = None
+    ):
+        if extra is None:
+            extra = {}
+        self._callback_extra = dict_to_struct(extra)
         with self._lock:
             callbacks = self.callbacks.get(control, {})
             for trigger in triggers:
@@ -71,8 +83,10 @@ def handle_task_result(task: asyncio.Task):
         task = asyncio.create_task(self._stream_events(), name=f"{viam._TASK_PREFIX}-input_stream_events")
         task.add_done_callback(handle_task_result)
 
-    async def trigger_event(self, event: Event, *, timeout: Optional[float] = None):
-        request = TriggerEventRequest(controller=self.name, event=event.proto)
+    async def trigger_event(self, event: Event, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None):
+        if extra is None:
+            extra = {}
+        request = TriggerEventRequest(controller=self.name, event=event.proto, extra=dict_to_struct(extra))
         try:
             await self.client.TriggerEvent(request, timeout=timeout)
         except GRPCError as e:
@@ -88,7 +102,7 @@ async def _stream_events(self):
         if not self.callbacks:
             return
 
-        request = StreamEventsRequest(controller=self.name, events=[])
+        request = StreamEventsRequest(controller=self.name, events=[], extra=self._callback_extra)
         with self._lock:
             for (control, callbacks) in self.callbacks.items():
                 event = StreamEventsRequest.Events(

src/viam/components/input/input.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from google.protobuf.timestamp_pb2 import Timestamp
 from typing_extensions import Self
@@ -136,7 +136,7 @@ class Controller(ComponentBase):
     """
 
     @abc.abstractmethod
-    async def get_controls(self, *, timeout: Optional[float] = None, **kwargs) -> List[Control]:
+    async def get_controls(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> List[Control]:
         """
         Returns a list of Controls provided by the Controller
 
@@ -146,7 +146,9 @@ async def get_controls(self, *, timeout: Optional[float] = None, **kwargs) -> Li
         ...
 
     @abc.abstractmethod
-    async def get_events(self, *, timeout: Optional[float] = None, **kwargs) -> Dict[Control, Event]:
+    async def get_events(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Dict[Control, Event]:
         """
         Returns the most recent Event for each input
         (which should be the current state)
@@ -157,7 +159,15 @@ async def get_events(self, *, timeout: Optional[float] = None, **kwargs) -> Dict
         ...
 
     @abc.abstractmethod
-    def register_control_callback(self, control: Control, triggers: List[EventType], function: Optional[ControlFunction], **kwargs):
+    def register_control_callback(
+        self,
+        control: Control,
+        triggers: List[EventType],
+        function: Optional[ControlFunction],
+        *,
+        extra: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
         """
         Register a function that will fire on given EventTypes for a given
         Control
@@ -171,7 +181,7 @@ def register_control_callback(self, control: Control, triggers: List[EventType],
         """
         ...
 
-    async def trigger_event(self, event: Event, *, timeout: Optional[float] = None, **kwargs):
+    async def trigger_event(self, event: Event, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs):
         """Directly send an Event (such as a button press) from external code
 
         Args:

src/viam/components/input/service.py

@@ -18,9 +18,10 @@
     TriggerEventRequest,
     TriggerEventResponse,
 )
+from viam.utils import struct_to_dict
 
-from .input import Control, Controller, Event, EventType
 
+from .input import Control, Controller, Event, EventType
 
 LOGGER = viam.logging.getLogger(__name__)
 
@@ -41,7 +42,7 @@ async def GetControls(self, stream: Stream[GetControlsRequest, GetControlsRespon
         except ComponentNotFoundError as e:
             raise e.grpc_error
         timeout = stream.deadline.time_remaining() if stream.deadline else None
-        controls = await controller.get_controls(timeout=timeout)
+        controls = await controller.get_controls(extra=struct_to_dict(request.extra), timeout=timeout)
         response = GetControlsResponse(controls=[c.value for c in controls])
         await stream.send_message(response)
 
@@ -54,7 +55,7 @@ async def GetEvents(self, stream: Stream[GetEventsRequest, GetEventsResponse]) -
         except ComponentNotFoundError as e:
             raise e.grpc_error
         timeout = stream.deadline.time_remaining() if stream.deadline else None
-        events = await controller.get_events(timeout=timeout)
+        events = await controller.get_events(extra=struct_to_dict(request.extra), timeout=timeout)
         pb_events = [e.proto for e in events.values()]
         response = GetEventsResponse(events=pb_events)
         await stream.send_message(response)
@@ -85,11 +86,21 @@ def ctrlFunc(event: Event):
         for event in request.events:
             triggers = [EventType(et) for et in event.events]
             if len(triggers):
-                controller.register_control_callback(Control(event.control), triggers, ctrlFunc)
+                controller.register_control_callback(
+                    Control(event.control),
+                    triggers,
+                    ctrlFunc,
+                    extra=struct_to_dict(request.extra),
+                )
 
             cancelled_triggers = [EventType(et) for et in event.cancelled_events]
             if len(cancelled_triggers):
-                controller.register_control_callback(Control(event.control), cancelled_triggers, None)
+                controller.register_control_callback(
+                    Control(event.control),
+                    cancelled_triggers,
+                    None,
+                    extra=struct_to_dict(request.extra),
+                )
 
         # Asynchronously wait for messages to come over the read pipe and run the READ function whenever the pipe is ready.
         def read():
@@ -130,7 +141,12 @@ def unregister_pipe_callbacks():
             for event in request.events:
                 triggers = [EventType(et) for et in event.events]
                 if len(triggers):
-                    controller.register_control_callback(Control(event.control), triggers, None)
+                    controller.register_control_callback(
+                        Control(event.control),
+                        triggers,
+                        None,
+                        extra=struct_to_dict(request.extra),
+                    )
 
     async def TriggerEvent(self, stream: Stream[TriggerEventRequest, TriggerEventResponse]) -> None:
         request = await stream.recv_message()
@@ -141,7 +157,7 @@ async def TriggerEvent(self, stream: Stream[TriggerEventRequest, TriggerEventRes
             controller = self.get_component(name)
             pb_event = request.event
             event = Event.from_proto(pb_event)
-            await controller.trigger_event(event, timeout=timeout)
+            await controller.trigger_event(event, extra=struct_to_dict(request.extra), timeout=timeout)
         except ComponentNotFoundError as e:
             raise e.grpc_error
         except NotSupportedError as e:

tests/mocks/components.py

@@ -473,8 +473,11 @@ def __init__(self, name: str):
         self.events: Dict[Control, Event] = {}
         self.callbacks: Dict[Control, Dict[EventType, Optional[ControlFunction]]] = {}
         self.timeout: Optional[float] = None
+        self.extra = None
+        self.reg_extra = None
 
-    async def get_controls(self, *, timeout: Optional[float] = None, **kwargs) -> List[Control]:
+    async def get_controls(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> List[Control]:
+        self.extra = extra
         self.timeout = timeout
         return [
             Control.ABSOLUTE_X,
@@ -500,15 +503,27 @@ async def get_controls(self, *, timeout: Optional[float] = None, **kwargs) -> Li
             Control.BUTTON_E_STOP,
         ]
 
-    async def get_events(self, *, timeout: Optional[float] = None, **kwargs) -> Dict[Control, Event]:
+    async def get_events(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Dict[Control, Event]:
+        self.extra = extra
         self.timeout = timeout
         return self.events
 
-    def register_control_callback(self, control: Control, triggers: List[EventType], function: Optional[ControlFunction], **kwargs):
+    def register_control_callback(
+        self,
+        control: Control,
+        triggers: List[EventType],
+        function: Optional[ControlFunction],
+        extra: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
         self.callbacks[control] = {trigger: function for trigger in triggers}
+        self.reg_extra = extra
 
-    async def trigger_event(self, event: Event, *, timeout: Optional[float] = None, **kwargs):
+    async def trigger_event(self, event: Event, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs):
         self.events[event.control] = event
+        self.extra = extra
         self.timeout = timeout
         callback = self.callbacks.get(event.control, {}).get(event.event)
         if callback: