Fix mypy protobuf in team com

Author: Flova

.gitignore

@@ -222,6 +222,7 @@ doku/*
 /bitbots_team_communication/bitbots_team_communication/bitbots_team_communication/RobocupProtocol
 # Protobuf generated file
 /bitbots_team_communication/bitbots_team_communication/bitbots_team_communication/robocup_extension_pb2.py
+/bitbots_team_communication/bitbots_team_communication/bitbots_team_communication/robocup_extension_pb2.pyi
 
 # Workspace git status file from the deploy tool
 **/workspace_status.json

bitbots_team_communication/bitbots_team_communication/CMakeLists.txt

@@ -6,19 +6,16 @@ find_package(ament_cmake_python REQUIRED)
 find_package(bitbots_docs REQUIRED)
 find_package(Protobuf REQUIRED)
 
-protobuf_generate_python(
-  PROTO_PY bitbots_team_communication/RobocupProtocol/robocup_extension.proto)
+add_custom_target(${PROJECT_NAME}_generate_proto ALL)
 
-add_custom_target(
-  bitbots_team_communication ALL
-  DEPENDS ${PROTO_PY}
-  COMMENT "Generating protobuf")
 add_custom_command(
-  TARGET bitbots_team_communication
-  POST_BUILD
-  COMMAND ${CMAKE_COMMAND} -E copy ${PROTO_PY}
-          ${CMAKE_SOURCE_DIR}/bitbots_team_communication
-  COMMENT "Copying protobuf to source dir")
+  TARGET ${PROJECT_NAME}_generate_proto
+  COMMAND protoc
+  --pyi_out ${CMAKE_SOURCE_DIR}/bitbots_team_communication/
+  --python_out ${CMAKE_SOURCE_DIR}/bitbots_team_communication/
+  --proto_path ${CMAKE_SOURCE_DIR}/bitbots_team_communication/RobocupProtocol
+  ${CMAKE_SOURCE_DIR}/bitbots_team_communication/RobocupProtocol/robocup_extension.proto
+  COMMENT "Generating protobuf python interface including stubs")
 
 enable_bitbots_docs()
 ament_python_install_package(${PROJECT_NAME})
@@ -50,7 +47,34 @@ if(BUILD_TESTING)
   endforeach()
 
   find_package(ament_cmake_mypy REQUIRED)
-  ament_mypy(CONFIG_FILE "${CMAKE_CURRENT_LIST_DIR}/mypy.ini")
+
+
+  # The following logic for manually filtering out the files can be removed if 
+  # https://github.com/ament/ament_lint/pull/516 is merged and released in our ROS distro
+  
+  # Collect all .py and .pyi files
+  file(GLOB_RECURSE PY_FILES "${CMAKE_CURRENT_LIST_DIR}/**/*.py")
+  file(GLOB_RECURSE PYI_FILES "${CMAKE_CURRENT_LIST_DIR}/**/*.pyi")
+
+  # Filter out .py files if a corresponding .pyi file exists
+  foreach(py_file IN LISTS PY_FILES)
+      # Get the directory and basename (without extension)
+      get_filename_component(py_dir "${py_file}" DIRECTORY)
+      get_filename_component(py_basename "${py_file}" NAME_WE)
+
+      # Construct the corresponding .pyi file path
+      set(pyi_file "${py_dir}/${py_basename}.pyi")
+      
+      # If a corresponding .pyi file exists, remove the .py file from the list
+      if(EXISTS "${pyi_file}")
+          list(REMOVE_ITEM PY_FILES "${py_file}")
+      endif()
+  endforeach()
+
+  ament_mypy(
+    CONFIG_FILE "${CMAKE_CURRENT_LIST_DIR}/mypy.ini"
+    ${PY_FILES} ${PYI_FILES}
+    )
 endif()
 
 ament_package()

bitbots_team_communication/bitbots_team_communication/bitbots_team_communication/converter/robocup_protocol_converter.py

@@ -16,40 +16,40 @@ class TeamColor(IntEnum):
 class RobocupProtocolConverter:
     def __init__(self, own_team_color: TeamColor):
         self.role_mapping = (
-            (Proto.Role.ROLE_UNDEFINED, Strategy.ROLE_UNDEFINED),
-            (Proto.Role.ROLE_IDLING, Strategy.ROLE_IDLING),
-            (Proto.Role.ROLE_OTHER, Strategy.ROLE_OTHER),
-            (Proto.Role.ROLE_STRIKER, Strategy.ROLE_STRIKER),
-            (Proto.Role.ROLE_SUPPORTER, Strategy.ROLE_SUPPORTER),
-            (Proto.Role.ROLE_DEFENDER, Strategy.ROLE_DEFENDER),
-            (Proto.Role.ROLE_GOALIE, Strategy.ROLE_GOALIE),
+            (Proto.ROLE_UNDEFINED, Strategy.ROLE_UNDEFINED),
+            (Proto.ROLE_IDLING, Strategy.ROLE_IDLING),
+            (Proto.ROLE_OTHER, Strategy.ROLE_OTHER),
+            (Proto.ROLE_STRIKER, Strategy.ROLE_STRIKER),
+            (Proto.ROLE_SUPPORTER, Strategy.ROLE_SUPPORTER),
+            (Proto.ROLE_DEFENDER, Strategy.ROLE_DEFENDER),
+            (Proto.ROLE_GOALIE, Strategy.ROLE_GOALIE),
         )
         self.action_mapping = (
-            (Proto.Action.ACTION_UNDEFINED, Strategy.ACTION_UNDEFINED),
-            (Proto.Action.ACTION_POSITIONING, Strategy.ACTION_POSITIONING),
-            (Proto.Action.ACTION_GOING_TO_BALL, Strategy.ACTION_GOING_TO_BALL),
-            (Proto.Action.ACTION_TRYING_TO_SCORE, Strategy.ACTION_TRYING_TO_SCORE),
-            (Proto.Action.ACTION_WAITING, Strategy.ACTION_WAITING),
-            (Proto.Action.ACTION_KICKING, Strategy.ACTION_KICKING),
-            (Proto.Action.ACTION_SEARCHING, Strategy.ACTION_SEARCHING),
-            (Proto.Action.ACTION_LOCALIZING, Strategy.ACTION_LOCALIZING),
+            (Proto.ACTION_UNDEFINED, Strategy.ACTION_UNDEFINED),
+            (Proto.ACTION_POSITIONING, Strategy.ACTION_POSITIONING),
+            (Proto.ACTION_GOING_TO_BALL, Strategy.ACTION_GOING_TO_BALL),
+            (Proto.ACTION_TRYING_TO_SCORE, Strategy.ACTION_TRYING_TO_SCORE),
+            (Proto.ACTION_WAITING, Strategy.ACTION_WAITING),
+            (Proto.ACTION_KICKING, Strategy.ACTION_KICKING),
+            (Proto.ACTION_SEARCHING, Strategy.ACTION_SEARCHING),
+            (Proto.ACTION_LOCALIZING, Strategy.ACTION_LOCALIZING),
         )
         self.side_mapping = (
-            (Proto.OffensiveSide.SIDE_UNDEFINED, Strategy.SIDE_UNDEFINED),
-            (Proto.OffensiveSide.SIDE_LEFT, Strategy.SIDE_LEFT),
-            (Proto.OffensiveSide.SIDE_MIDDLE, Strategy.SIDE_MIDDLE),
-            (Proto.OffensiveSide.SIDE_RIGHT, Strategy.SIDE_RIGHT),
+            (Proto.SIDE_UNDEFINED, Strategy.SIDE_UNDEFINED),
+            (Proto.SIDE_LEFT, Strategy.SIDE_LEFT),
+            (Proto.SIDE_MIDDLE, Strategy.SIDE_MIDDLE),
+            (Proto.SIDE_RIGHT, Strategy.SIDE_RIGHT),
         )
 
         self.proto_to_team_data_team_mapping = {
-            Proto.Team.UNKNOWN_TEAM: RobotRelative.ROBOT_UNDEFINED,
-            Proto.Team.BLUE: RobotRelative.ROBOT_BLUE,
-            Proto.Team.RED: RobotRelative.ROBOT_RED,
+            Proto.UNKNOWN_TEAM: RobotRelative.ROBOT_UNDEFINED,
+            Proto.BLUE: RobotRelative.ROBOT_BLUE,
+            Proto.RED: RobotRelative.ROBOT_RED,
         }
         self.state_to_proto_team_mapping = {
-            RobotAttributes.TEAM_OWN: Proto.Team.RED if own_team_color == TeamColor.RED else Proto.Team.BLUE,
-            RobotAttributes.TEAM_OPPONENT: Proto.Team.BLUE if own_team_color == TeamColor.RED else Proto.Team.RED,
-            RobotAttributes.TEAM_UNKNOWN: Proto.Team.UNKNOWN_TEAM,
+            RobotAttributes.TEAM_OWN: Proto.RED if own_team_color == TeamColor.RED else Proto.BLUE,
+            RobotAttributes.TEAM_OPPONENT: Proto.BLUE if own_team_color == TeamColor.RED else Proto.RED,
+            RobotAttributes.TEAM_UNKNOWN: Proto.UNKNOWN_TEAM,
         }
 
         proto_to_team_data_mappings = {

bitbots_team_communication/bitbots_team_communication/bitbots_team_communication/converter/state_to_message_converter.py

@@ -25,9 +25,9 @@ def convert(
     ) -> Proto.Message:
         def convert_gamestate(gamestate: Optional[GameState], message: Proto.Message):
             if gamestate is not None and is_still_valid_checker(gamestate.header.stamp):
-                message.state = Proto.State.PENALISED if gamestate.penalized else Proto.State.UNPENALISED
+                message.state = Proto.PENALISED if gamestate.penalized else Proto.UNPENALISED
             else:
-                message.state = Proto.State.UNKNOWN_STATE
+                message.state = Proto.UNKNOWN_STATE
 
             return message
 

bitbots_team_communication/bitbots_team_communication/mypy.ini

@@ -3,4 +3,6 @@
 [mypy]
 check_untyped_defs = True
 ignore_missing_imports = True
-exclude = .*_pb2\.py
+
+# Ignore var-annotations because it is in autogenerated code from protobuf
+disable_error_code = var-annotated

bitbots_team_communication/bitbots_team_communication/test/converter/test_message_to_team_data_converter.py

@@ -60,9 +60,9 @@ def convert_from_message(message: Proto.Message, team_data: TeamData | None = No
 
 @pytest.fixture
 def message_with_strategy(message) -> Proto.Message:
-    message.role = Proto.Role.ROLE_STRIKER
-    message.action = Proto.Action.ACTION_KICKING
-    message.offensive_side = Proto.OffensiveSide.SIDE_RIGHT
+    message.role = Proto.ROLE_STRIKER
+    message.action = Proto.ACTION_KICKING
+    message.offensive_side = Proto.SIDE_RIGHT
 
     return message
 
@@ -111,9 +111,9 @@ def message() -> Proto.Message:
 def robot_message(player_id, is_own_team):
     robot = Proto.Robot()
     robot.player_id = player_id
-    robot.team = own_team_id
+    robot.team = own_team_id  # type: ignore[assignment]
     if not is_own_team:
-        robot.team = own_team_id + 1
+        robot.team = own_team_id + 1  # type: ignore[assignment]
 
     set_position(robot.position)
     set_covariance_matrix(robot.covariance)

bitbots_team_communication/bitbots_team_communication/test/converter/test_robocup_protocol_converter.py

@@ -20,9 +20,9 @@ def test_setup_of_mappings(snapshot):
 
 def test_setup_of_team_color_mapping(snapshot):
     converter = protocol_converter()
-    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_OWN] == Proto.Team.BLUE
-    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_OPPONENT] == Proto.Team.RED
-    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_UNKNOWN] == Proto.Team.UNKNOWN_TEAM
+    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_OWN] == Proto.BLUE
+    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_OPPONENT] == Proto.RED
+    assert converter.state_to_proto_team_mapping[RobotAttributes.TEAM_UNKNOWN] == Proto.UNKNOWN_TEAM
     assert converter.proto_to_team_data_team_mapping == snapshot
 
 

bitbots_team_communication/bitbots_team_communication/test/converter/test_state_to_message_converter.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from unittest.mock import Mock
 
 import numpy
@@ -44,23 +45,23 @@ def test_convert_gamestate(state_with_gamestate):
 
     gamestate.penalized = False
     result = convert_to_message(state_with_gamestate)
-    assert result.state == Proto.State.UNPENALISED
+    assert result.state == Proto.UNPENALISED
 
     gamestate.penalized = True
     result = convert_to_message(state_with_gamestate)
-    assert result.state == Proto.State.PENALISED
+    assert result.state == Proto.PENALISED
 
     assert validity_checker_valid.call_count == 2
     validity_checker_valid.assert_called_with(gamestate.header.stamp)
 
 
 def test_convert_gamestate_expired_headers(state_with_gamestate):
     message = Proto.Message()
-    message.state = Proto.State.UNPENALISED
+    message.state = Proto.UNPENALISED
 
     result = convert_to_message(state_with_gamestate, message, is_state_expired=True)
 
-    assert result.state == Proto.State.UNKNOWN_STATE
+    assert result.state == Proto.UNKNOWN_STATE
 
 
 def test_convert_current_pose(snapshot, state_with_current_pose):
@@ -164,22 +165,22 @@ def test_convert_obstacles_to_robots_expired_headers(state_with_seen_robots):
 def test_convert_strategy(state_with_strategy):
     result = convert_to_message(state_with_strategy)
 
-    assert result.role == Proto.Role.ROLE_STRIKER
-    assert result.action == Proto.Action.ACTION_KICKING
-    assert result.offensive_side == Proto.OffensiveSide.SIDE_RIGHT
+    assert result.role == Proto.ROLE_STRIKER
+    assert result.action == Proto.ACTION_KICKING
+    assert result.offensive_side == Proto.SIDE_RIGHT
 
     validity_checker_valid.assert_called_with(state_with_strategy.strategy_time)
 
 
 def test_convert_strategy_expired_headers(state_with_strategy):
     message = Proto.Message()
-    message.role = Proto.Role.ROLE_DEFENDER
+    message.role = Proto.ROLE_DEFENDER
 
     result = convert_to_message(state_with_strategy, message, is_state_expired=True)
 
-    assert result.role == Proto.Role.ROLE_DEFENDER
-    assert result.action == Proto.Action.ACTION_UNDEFINED
-    assert result.offensive_side == Proto.OffensiveSide.SIDE_UNDEFINED
+    assert result.role == Proto.ROLE_DEFENDER
+    assert result.action == Proto.ACTION_UNDEFINED
+    assert result.offensive_side == Proto.SIDE_UNDEFINED
 
 
 def test_convert_time_to_ball(state_with_time_to_ball):
@@ -208,7 +209,7 @@ def test_convert_time_to_ball_expired_headers(state_with_time_to_ball):
     assert pytest.approx(result.time_to_ball) == 9999.0
 
 
-def convert_to_message(team_data_state, message: Proto.Message = None, is_state_expired=False):
+def convert_to_message(team_data_state, message: Optional[Proto.Message] = None, is_state_expired=False):
     message = message if message else Proto.Message()
     validity_checker = validity_checker_expired if is_state_expired else validity_checker_valid
     return RobocupProtocolConverter(own_team_color).convert_to_message(team_data_state, message, validity_checker)