sbhs_robomaster.client

  1from enum import Enum
  2from types import TracebackType
  3from typing import Any, Coroutine, Final, Optional, Type
  4import asyncio
  5from time import time
  6from .data import *
  7from .push_receiver import PushReceiver
  8from .feed import Feed
  9from .auto_connect import AutoConnectProtocol
 10
 11
 12DIRECT_CONNECT_IP: Final[str] = "192.168.2.1"
 13IP_PORT: Final[int] = 40926
 14
 15_CONTROL_PORT: Final[int] = 40923
 16_PUSH_PORT: Final[int] = 40924
 17
 18
 19async def auto_connect_to_robomaster() -> "RoboMasterClient":
 20    """
 21    It is preferable to use `connect_to_robomaster` where possible, especially if
 22    you are directly connected to the WiFi the robot provides (use `DIRECT_CONNECT_IP`
 23    in that case.) This is due to it not accidentally getting confused when there
 24    are multiple robots (see below).
 25
 26    Automatically finds an unpaired RoboMaster robot on the network and connects to it.
 27    
 28    This assumes there will only be ONE unpaired RoboMaster robot on the network, or you
 29    may end up connecting to a random robot.
 30
 31    ```py
 32    import asyncio
 33    from sbhs_robomaster import auto_connect_to_robomaster
 34
 35    async def main():
 36        async with await auto_connect_to_robomaster() as robot:
 37            # Do stuff with robot
 38            pass
 39
 40    asyncio.run(main())
 41    ```
 42    """
 43
 44    loop = asyncio.get_event_loop()
 45    transport, protocol = await loop.create_datagram_endpoint(AutoConnectProtocol, ("", IP_PORT))
 46
 47    ip = await protocol.ip_future
 48
 49    transport.close()
 50
 51    return await connect_to_robomaster(ip)
 52
 53
 54async def connect_to_robomaster(ip: str) -> "RoboMasterClient":
 55    """
 56    Connects to a RoboMaster robot at the given IP address.
 57    
 58    If the host computer (the computer the code is running on) is connected directly to robot's WiFi,
 59    and not over another network, pass in `DIRECT_CONNECT_IP`:
 60
 61    ```py
 62    import asyncio
 63    from sbhs_robomaster import connect_to_robomaster, DIRECT_CONNECT_IP
 64
 65    async def main():
 66        async with await connect_to_robomaster(DIRECT_CONNECT_IP) as robot:
 67            # Do stuff with robot
 68            pass
 69
 70    asyncio.run(main())
 71    ```
 72    """
 73
 74    command_socket = await asyncio.open_connection(ip, _CONTROL_PORT)
 75
 76    client = RoboMasterClient(command_socket)
 77    await client.async_init()
 78
 79    return client
 80
 81
 82def command_to_str(command: str | int | float | bool | Enum) -> str:
 83    """
 84    Converts a data type to a string that can be sent to the robot.
 85    """
 86
 87    match command:
 88        case str():
 89            return command
 90        case bool():
 91            return "on" if command else "off"
 92        case int() | float():
 93            return str(command)
 94        case Enum():
 95            if isinstance(command.value, str) or isinstance(command.value, int) or isinstance(command.value, float) \
 96                or isinstance(command.value, bool) or isinstance(command.value, Enum):
 97                return command_to_str(command.value)
 98            else:
 99                raise Exception("Invalid enum type")
100
101
102class RoboMasterClient:
103    line: Feed[Line]
104    """
105    A feed of line recognition data.
106
107    Line recognition has to be enabled first with `RoboMasterClient.set_line_recognition_enabled` and
108    the line colour has to be set with `RoboMasterClient.set_line_recognition_colour`.
109    """
110
111    rotation: Feed[ChassisRotation]
112    """
113    A feed of the robot's rotation.
114
115    This has to be enabled first with `RoboMasterClient.set_rotation_push_rate`.
116    """
117
118    attitude: Feed[ChassisAttitude]
119    """
120    A feed of the robot's attitude.
121
122    This has to be enabled first with `RoboMasterClient.set_attitude_push_rate`.
123    """
124
125    status: Feed[ChassisStatus]
126    """
127    A feed of the robot's status.
128
129    This has to be enabled first with `RoboMasterClient.set_status_push_rate`.
130    """
131
132    _line_enabled: bool
133    _rotation_frequency: Frequency
134    _attitude_frequency: Frequency
135    _status_frequency: Frequency
136
137    _conn: tuple[asyncio.StreamReader, asyncio.StreamWriter]
138    _command_lock: asyncio.Lock
139    _push_receiver: PushReceiver
140    _handle_push_task: asyncio.Task[None]
141    _exiting: bool
142
143    def __init__(self, command_conn: tuple[asyncio.StreamReader, asyncio.StreamWriter]) -> None:
144        """
145        This constructor shouldn't be used directly. Use `connect_to_robomaster` instead.
146        """
147
148        self.line = Feed()
149        self.rotation = Feed()
150        self.attitude = Feed()
151        self.status = Feed()
152
153        self._line_enabled = False
154        self._rotation_frequency = Frequency.Off
155        self._attitude_frequency = Frequency.Off
156        self._status_frequency = Frequency.Off
157
158        self._conn = command_conn
159        self._command_lock = asyncio.Lock()
160        self._push_receiver = PushReceiver(_PUSH_PORT)
161        self._handle_push_task = asyncio.create_task(self._handle_push(self._push_receiver.feed))
162        self._exiting = False
163
164    async def async_init(self):
165        """
166        Performs any asynchronous initialisation required. Should be called
167        immediately after the class is constructed. `connect_to_robomaster`
168        handles this automatically.
169        """
170        await self.do("command")
171
172        await self.set_line_recognition_enabled(False)
173        await self.set_rotation_push_rate(Frequency.Off)
174        await self.set_attitude_push_rate(Frequency.Off)
175        await self.set_status_push_rate(Frequency.Off)
176
177    async def __aenter__(self):
178        return self
179
180    async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], tb: Optional[TracebackType]):
181        await self.do("quit")
182
183        self._exiting = True
184
185        self._handle_push_task.cancel()
186        self._conn[1].close()
187
188    async def _handle_push(self, feed: Feed[Response]):
189        while True:
190            response = await feed.get()
191
192            topic = response.data[0]
193            subject = response.data[2]
194            parseData = Response(response.data[3:])
195
196            if topic == "chassis":
197                if subject == "position":   self.rotation.feed(ChassisRotation.parse(parseData))
198                elif subject == "attitude": self.attitude.feed(ChassisAttitude.parse(parseData))
199                elif subject == "status":   self.status.feed(ChassisStatus.parse(parseData))
200                else:                       print(f"Unknown chassis subject: {subject}")
201            elif topic == "AI":
202                if subject == "line": self.line.feed(Line.parse(parseData))
203                else:                 print(f"Unknown AI subject: {subject}")
204            else:
205                print(f"Unknown topic: {topic}")
206
207    async def _do(self, command: str):
208        async with self._command_lock:
209            if self._exiting:
210                raise Exception("Client is exiting.")
211
212            self._conn[1].write(command.encode())
213            await self._conn[1].drain()
214
215            data = await self._conn[0].readuntil(b";")
216
217            return Response(data.decode()[:-1].split(" "))
218
219    async def do(self, *command_data: str | int | float | bool | Enum) -> Response:
220        """
221        Low level command sending. Every function that influences the robot's behaviour
222        uses this internally.
223
224        It's better to use the higher level functions instead of this.
225
226        Arguments:
227         - *command_data: The parts of the command to send to the robot. These can be strings, numbers, booleans or enums.
228                          All of these will be converted to strings using `command_to_str`.
229        """
230        
231        command = ' '.join(map(command_to_str, command_data)) + ';'
232        return await self._do(command)
233
234    async def get_version(self) -> str:
235        return (await self.do("version")).get_str(0)
236
237    async def set_speed(self, forwards: float, right: float, clockwise: float = 0) -> None:
238        """
239        Causes the robot to move with the given speeds indefinitely.
240
241        Arguments:
242         - forwards: The forwards speed in m/s
243         - right: The right speed in m/s
244         - clockwise: The rotational speed in degrees/s clockwise
245        """
246
247        assert -3.5 <= forwards <= 3.5, "Movement speed must be between -3.5 and 3.5 m/s"
248        assert -3.5 <= right <= 3.5,    "Movement speed must be between -3.5 and 3.5 m/s"
249        assert -600 <= clockwise <= 600, "Rotation speed must be between -600 and 600 degrees/s"
250
251        await self.do("chassis", "speed", "x", forwards, "y", right, "z", clockwise)
252
253    async def set_wheel_speed(self, front_right: float, front_left: float, back_left: float, back_right: float) -> None:
254        """
255        All speeds are in rpm.
256        """
257
258        assert -1000 <= front_right <= 1000, "Wheel speeds must be between -1000 and 1000 rpm"
259        assert -1000 <= front_left <= 1000,  "Wheel speeds must be between -1000 and 1000 rpm"
260        assert -1000 <= back_left <= 1000,   "Wheel speeds must be between -1000 and 1000 rpm"
261        assert -1000 <= back_right <= 1000,  "Wheel speeds must be between -1000 and 1000 rpm"
262
263        await self.do(
264            "chassis", "wheel",
265            "w1", front_right,
266            "w2", front_left,
267            "w3", back_left,
268            "w4", back_right
269        )
270
271    async def set_left_right_wheel_speeds(self, left: float, right: float) -> None:
272        """
273        All speeds are in rpm.
274        """
275        await self.set_wheel_speed(right, left, left, right)
276
277    async def set_all_wheel_speeds(self, speed: float) -> None:
278        """
279        All speeds are in rpm.
280        """
281        await self.set_wheel_speed(speed, speed, speed, speed)
282
283    async def get_speed(self) -> ChassisSpeed:
284        return ChassisSpeed.parse(await self.do("chassis", "speed", "?"))
285
286    async def rotate(self, clockwise: float, rotation_speed: float | None = None, timeout: float = 10) -> None:
287        """
288        **NOTE:** This function constantly polls the robot for its rotation and
289        may starve out other areas in the program trying to send messages to
290        the robot. In general, due to the serial nature of the communication
291        protocol, any communication with the robot should not be largely
292        parallel.
293
294        Arguments:
295         - clockwise: The degrees clockwise to rotate.
296         - rotation_speed: The speed to rotate, in degrees/s. (Must be positive.)
297         - timeout: How long to wait for the turn to finish (in seconds).
298        """
299
300        assert -1800 <= clockwise <= 1800, "Rotation must be between -1800 and 1800 degrees"
301        
302        if rotation_speed is not None:
303            assert 0 <= rotation_speed <= 600, "Rotation speed must be positive and less than 600 degrees/s"
304
305        args = [
306            "chassis", "move",
307            "x", 0,
308            "y", 0,
309            "z", clockwise
310        ]
311
312        if rotation_speed is not None:
313            args.append("vz")
314            args.append(rotation_speed)
315
316        await self.do(*args)
317
318        start_time = time()
319        while time() - start_time <= timeout:
320            status = await self.get_status()
321
322            if status.static:
323                return
324            
325            await asyncio.sleep(0.02) # 50Hz
326            
327        raise TimeoutError()
328
329    async def get_rotation(self) -> ChassisRotation:
330        return ChassisRotation.parse(await self.do("chassis", "position", "?"))
331    
332    async def get_attitude(self) -> ChassisAttitude:
333        return ChassisAttitude.parse(await self.do("chassis", "attitude", "?"))
334    
335    async def get_status(self) -> ChassisStatus:
336        return ChassisStatus.parse(await self.do("chassis", "status", "?"))
337
338    async def set_rotation_push_rate(self, freq: Frequency) -> None:
339        """
340        Sets the push rate of the chassis rotation in Hz. Use `Frequency.Off` to disable.
341
342        Access the rotation data through the `RoboMasterClient.rotation` feed.
343        """
344
345        # This sets the "position" push frequency because
346        # that is where the rotation is stored.
347        if freq == Frequency.Off:
348            await self.do("chassis", "push", "position", False)
349        else:
350            await self.do("chassis", "push", "position", True, "pfreq", freq)
351
352        self._rotation_frequency = freq
353
354    def get_rotation_push_rate(self) -> Frequency:
355        """
356        Gets the push rate of the chassis rotation in Hz. `Frequency.Off` means it is disabled.
357        """
358        return self._rotation_frequency
359
360    async def set_attitude_push_rate(self, freq: Frequency) -> None:
361        """
362        Sets the push rate of the chassis attitude in Hz. Use `Frequency.Off` to disable.
363
364        Access the attitude data through the `RoboMasterClient.attitude` feed.
365        """
366
367        if freq == Frequency.Off:
368            await self.do("chassis", "push", "attitude", False)
369        else:
370            await self.do("chassis", "push", "attitude", True, "afreq", freq)
371
372        self._attitude_frequency = freq
373
374    def get_attitude_push_rate(self) -> Frequency:
375        """
376        Gets the push rate of the chassis rotation in Hz. `Frequency.Off` means it is disabled.
377        """
378        return self._attitude_frequency
379
380    async def set_status_push_rate(self, freq: Frequency) -> None:
381        """
382        Sets the push rate of the chassis status in Hz. Use `Frequency.Off` to disable.
383
384        Access the status data through the `RoboMasterClient.status` feed.
385        """
386
387        if freq == Frequency.Off:
388            await self.do("chassis", "push", "status", False)
389        else:
390            await self.do("chassis", "push", "status", True, "sfreq", freq)
391
392        self._status_frequency = freq
393
394    def get_status_push_rate(self) -> Frequency:
395        """
396        Gets the push rate of the chassis rotation in Hz. `Frequency.Off` means it is disabled.
397        """
398        return self._status_frequency
399
400    async def set_ir_enabled(self, enabled: bool = True) -> None:
401        await self.do("ir_distance_sensor", "measure", enabled)
402    
403    async def get_ir_distance(self, ir_id: int) -> float:
404        """
405        The IR sensor has to be enabled first with `set_ir_enabled`.
406
407        Arguments:
408         - ir_id: The ID of the IR sensor. There appears to be only one IR sensor, so this should always be `1`.
409
410        Returns:
411         - The distance in centimeters.
412        """
413        return (await self.do("ir_distance_sensor", "distance", ir_id, "?")).get_float(0)
414    
415    async def move_arm(self, x_dist: float, y_dist: float) -> None:
416        """
417        Moves the robotic arm by the given distance relative to its current position.
418        To move to a specific position, use `set_arm_position` instead.
419
420        The units are unknown. Physically move the arm around and record the results from
421        `get_arm_position` to determine the desired inputs for this function.
422        """
423        await self.do("robotic_arm", "move", "x", x_dist, "y", y_dist)
424    
425    async def set_arm_position(self, x: float, y: float) -> None:
426        """
427        Moves the robotic arm to the given position. To move relative to the current position,
428        use `move_arm` instead.
429
430        The units are unknown. Physically move the arm to the desired position and then use
431        `get_arm_position` to find the desired inputs for this function. The origin for the
432        x and y values appears to be different for every robot, so values for one robot will
433        not carry over to another.
434        """
435        await self.do("robotic_arm", "moveto", "x", x, "y", y)
436    
437    async def get_arm_position(self) -> tuple[float, float]:
438        """
439        Returns:
440         - The current position of the robotic arm in unknown units. The format is `(x, y)`.
441        """
442
443        response = await self.do("robotic_arm", "position", "?")
444        return (response.get_float(0), response.get_float(1))
445    
446    # TODO Allow different levels of force when opening?
447    # https://robomaster-dev.readthedocs.io/en/latest/text_sdk/protocol_api.html#mechanical-gripper-opening-control
448    async def open_gripper(self) -> None:
449        await self.do("robotic_gripper", "open", 1)
450    
451    # TODO Allow different levels of force when closing?
452    # https://robomaster-dev.readthedocs.io/en/latest/text_sdk/protocol_api.html#mechanical-gripper-opening-control
453    async def close_gripper(self) -> None:
454        """
455        This closes the gripper until it encounters any resistance. If there is an object in
456        the way, the gripper will close until it is tightly gripping the object, but it will
457        not damage itself.
458        """
459        await self.do("robotic_gripper", "close", 1)
460
461    async def get_gripper_status(self) -> GripperStatus:
462        """
463        Gets whether the gripper is Open, Closed, or Partially Open.
464        """
465        return (await self.do("robotic_gripper", "status", "?")).get_enum(0, GripperStatus)
466
467    async def set_line_recognition_colour(self, colour: LineColour) -> None:
468        await self.do("AI", "attribute", "line_color", colour)
469
470    async def set_line_recognition_enabled(self, enabled: bool = True) -> None:
471        """
472        Access the line data through the `RoboMasterClient.line` feed.
473        """
474        await self.do("AI", "push", "line", enabled)
475
476        self._line_enabled = enabled
477
478    def get_line_recognition_enabled(self) -> bool:
479        return self._line_enabled
480    
481    async def set_leds(self, leds: LedPosition, colour: Colour, effect: LedEffect = LedEffect.Solid) -> None:
482        """
483        Sets the given LEDs to display a colour with a certain effect. Use
484        `.data.LedEffect.Off` to turn the leds off.
485
486        See `.data.LedPosition` to see how to pass in multiple leds at once.
487        
488        See `.data.LedEffect` to see the available effects.
489        """
490
491        async def set_led(led: str):
492            await self.do("led", "control", "comp", led, "r", colour.r, "g", colour.g, "b", colour.b, "effect", effect)
493
494        if leds == LedPosition.All:
495            await set_led("bottom_all")
496        else:
497            tasks: list[Coroutine[Any, Any, None]] = []
498
499            if LedPosition.Front in leds:
500                tasks.append(set_led("bottom_front"))
501            
502            if LedPosition.Back in leds:
503                tasks.append(set_led("bottom_back"))
504
505            if LedPosition.Left in leds:
506                tasks.append(set_led("bottom_left"))
507
508            if LedPosition.Right in leds:
509                tasks.append(set_led("bottom_right"))
510
511            await asyncio.gather(*tasks)