Merge pull request #31 from RLBot/context_managers

Add ContextManager pattern to MatchManager and RenderingManager plus default color

Author: VirxEC

rlbot/managers/match.py

@@ -33,6 +33,12 @@ def __init__(
         self.rlbot_interface: SocketRelay = SocketRelay("")
         self.rlbot_interface.packet_handlers.append(self._packet_reporter)
 
+    def __enter__(self) -> 'MatchManager':
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.disconnect()
+
     def ensure_server_started(self):
         """
         Ensures that RLBotServer is running, starting it if it is not.

rlbot/managers/rendering.py

@@ -1,5 +1,6 @@
 import math
 from collections.abc import Callable, Sequence
+from contextlib import contextmanager
 from typing import Optional
 
 from rlbot import flat
@@ -51,6 +52,8 @@ class Renderer:
     _group_id: Optional[int] = None
     _current_renders: list[flat.RenderMessage] = []
 
+    _default_color = white
+
     _screen_width_factor = 1.0
     _screen_height_factor = 1.0
 
@@ -72,6 +75,12 @@ def set_resolution(self, screen_width: float, screen_height: float):
         self._screen_width_factor = 1.0 / screen_width
         self._screen_height_factor = 1.0 / screen_height
 
+    def set_default_color(self, color: flat.Color):
+        """
+        Set which color to use when no other color is provided.
+        """
+        self._default_color = color
+
     @staticmethod
     def create_color(red: int, green: int, blue: int, alpha: int = 255) -> flat.Color:
         return flat.Color(red, green, blue, alpha)
@@ -111,6 +120,30 @@ def team_color(team: int, alt_color: bool = False) -> flat.Color:
     def _get_group_id(group_id: str) -> int:
         return hash(str(group_id).encode("utf-8")) % MAX_INT
 
+    @contextmanager
+    def context(self, group_id: str=DEFAULT_GROUP_ID, default_color=None):
+        """
+        Starts rendering as a context usable in with-statements.
+        After the with-statement the rendering is automatically ended.
+        Note, the is not possible to have two nested renderings started.
+
+        Example:
+
+        ```
+        with renderer.context(default_color=renderer.red):
+            renderer.draw_line_3d(car.pos, ball.pos)
+            renderer.draw_line_3d(car.pos, goal.pos)
+            renderer.draw_line_3d(ball.pos, goal.pos)
+        ```
+        """
+        try:
+            self.begin_rendering(group_id)
+            if default_color:
+                self.set_default_color(default_color)
+            yield
+        finally:
+            self.end_rendering()
+
     def begin_rendering(self, group_id: str = DEFAULT_GROUP_ID):
         """
         Begins a new render group. All render messages added after this call will be part of this group.
@@ -189,29 +222,29 @@ def draw_line_3d(
         self,
         start: flat.RenderAnchor | flat.BallAnchor | flat.CarAnchor | flat.Vector3,
         end: flat.RenderAnchor | flat.BallAnchor | flat.CarAnchor | flat.Vector3,
-        color: flat.Color,
+        color: flat.Color | None = None,
     ):
         """
         Draws a line between two anchors in 3d space.
         """
-        self.draw(flat.Line3D(_get_anchor(start), _get_anchor(end), color))
+        self.draw(flat.Line3D(_get_anchor(start), _get_anchor(end), color or self._default_color))
 
     def draw_polyline_3d(
         self,
         points: Sequence[flat.Vector3],
-        color: flat.Color,
+        color: flat.Color | None = None,
     ):
         """
         Draws a line going through each of the provided points.
         """
-        self.draw(flat.PolyLine3D(points, color))
+        self.draw(flat.PolyLine3D(points, color or self._default_color))
 
     def draw_string_3d(
         self,
         text: str,
         anchor: flat.RenderAnchor | flat.BallAnchor | flat.CarAnchor | flat.Vector3,
         scale: float,
-        foreground: flat.Color,
+        foreground: flat.Color | None = None,
         background: flat.Color = flat.Color(a=0),
         h_align: flat.TextHAlign = flat.TextHAlign.Left,
         v_align: flat.TextVAlign = flat.TextVAlign.Top,
@@ -225,7 +258,7 @@ def draw_string_3d(
                 text,
                 _get_anchor(anchor),
                 scale,
-                foreground,
+                foreground or self._default_color,
                 background,
                 h_align,
                 v_align,
@@ -238,7 +271,7 @@ def draw_string_2d(
         x: float,
         y: float,
         scale: float,
-        foreground: flat.Color,
+        foreground: flat.Color | None = None,
         background: flat.Color = flat.Color(a=0),
         h_align: flat.TextHAlign = flat.TextHAlign.Left,
         v_align: flat.TextVAlign = flat.TextVAlign.Top,
@@ -255,7 +288,7 @@ def draw_string_2d(
                 x * self._screen_width_factor,
                 y * self._screen_height_factor,
                 scale,
-                foreground,
+                foreground or self._default_color,
                 background,
                 h_align,
                 v_align,
@@ -268,7 +301,7 @@ def draw_rect_2d(
         y: float,
         width: float,
         height: float,
-        color: flat.Color,
+        color: flat.Color | None = None,
         h_align: flat.TextHAlign = flat.TextHAlign.Left,
         v_align: flat.TextVAlign = flat.TextVAlign.Top,
     ):
@@ -284,7 +317,7 @@ def draw_rect_2d(
                 y * self._screen_height_factor,
                 width * self._screen_width_factor,
                 height * self._screen_height_factor,
-                color,
+                color or self._default_color,
                 h_align,
                 v_align,
             )
@@ -295,7 +328,7 @@ def draw_rect_3d(
         anchor: flat.RenderAnchor | flat.BallAnchor | flat.CarAnchor | flat.Vector3,
         width: float,
         height: float,
-        color: flat.Color,
+        color: flat.Color | None = None,
         h_align: flat.TextHAlign = flat.TextHAlign.Left,
         v_align: flat.TextVAlign = flat.TextVAlign.Top,
     ):
@@ -310,7 +343,7 @@ def draw_rect_3d(
                 _get_anchor(anchor),
                 width * self._screen_width_factor,
                 height * self._screen_height_factor,
-                color,
+                color or self._default_color,
                 h_align,
                 v_align,
             )

rlbot/version.py

@@ -1 +1 @@
-__version__ = "2.0.0-beta.43"
+__version__ = "2.0.0-beta.44"

tests/render_test/render.py

@@ -1,3 +1,5 @@
+from email.policy import default
+
 from rlbot import flat
 from rlbot.flat import BallAnchor, CarAnchor, Color, RenderAnchor, Vector3
 from rlbot.managers import Script
@@ -30,12 +32,12 @@ def handle_packet(self, packet: flat.GamePacket):
 
             self.do_render(radius)
 
-        self.renderer.begin_rendering('tick')
-        hsv = self.renderer.create_color_hsv(packet.match_info.seconds_elapsed * 0.1, 1.0, 1.0)
-        self.renderer.set_resolution(1920, 1080)
-        self.renderer.draw_string_2d('HSV 300px 50px', 300, 50, 1.0, hsv)
-        self.renderer.set_resolution(1, 1)
-        self.renderer.end_rendering()
+        with self.renderer.context('tick', self.renderer.red):
+            self.renderer.set_resolution(1920, 1080)
+            hsv = self.renderer.create_color_hsv(packet.match_info.seconds_elapsed * 0.1, 1.0, 1.0)
+            self.renderer.draw_string_2d('HSV 300px 50px', 300, 50, 1.0, hsv)
+            self.renderer.draw_string_2d('Red 330px 70px', 330, 70, 1.0)  # Use default color
+            self.renderer.set_resolution(1, 1)
 
     def do_render(self, radius: float):
         self.renderer.begin_rendering()
@@ -75,7 +77,7 @@ def do_render(self, radius: float):
         )
 
         self.renderer.draw_rect_2d(
-            0.75, 0.75, 0.1, 0.1, Color(150, 30, 100), centered=False
+            0.75, 0.75, 0.1, 0.1, Color(150, 30, 100)
         )
         self.renderer.draw_rect_2d(0.75, 0.75, 0.1, 0.1, self.renderer.black)
         for hkey, h in {

tests/run_match.py

@@ -10,14 +10,10 @@
 RLBOT_SERVER_FOLDER = DIR / "../../core/RLBotCS/bin/Release/"
 
 if __name__ == "__main__":
-    match_manager = MatchManager(RLBOT_SERVER_FOLDER)
+    with MatchManager(RLBOT_SERVER_FOLDER) as man:
+        man.start_match(MATCH_CONFIG_PATH)
+        assert man.packet is not None
 
-    match_manager.start_match(MATCH_CONFIG_PATH)
-    assert match_manager.packet is not None
-
-    try:
         # wait for the match to end
-        while match_manager.packet.match_info.match_phase != flat.MatchPhase.Ended:
+        while man.packet.match_info.match_phase != flat.MatchPhase.Ended:
             sleep(1.0)
-    finally:
-        match_manager.shut_down()

tests/run_only.py

@@ -14,13 +14,11 @@
         match_config_path = Path(sys.argv[1])
     assert match_config_path.exists(), f"Match config not found: {match_config_path}"
 
-    # start the match
-    match_manager = MatchManager(RLBOT_SERVER_FOLDER)
-    match_manager.start_match(match_config_path, False)
+    with MatchManager(RLBOT_SERVER_FOLDER) as man:
+        man.start_match(match_config_path, False)
 
-    # wait
-    input("\nPress enter to end the match: ")
+        # Wait for input
+        input("\nPress enter to end the match: ")
 
-    # end the match and disconnect
-    match_manager.stop_match()
-    match_manager.disconnect()
+        # End the match
+        man.stop_match()