Resolve string-based type annotations in message handlers.

When `from __future__ import annotations` is active or forward references are used as strings, type annotations are not directly class objects. This change uses `get_type_hints` to resolve these string annotations to their actual types, ensuring correct validation of message handler signatures.

PiperOrigin-RevId: 942608060
Change-Id: Ic2b214ee85e835616cc3a18ec2699eba6deebb6a
This commit is contained in:
Matija Kecman
2026-07-04 16:02:54 -07:00
committed by Copybara-Service
parent 4a9364082e
commit 4cb14bef47
+18 -1
View File
@@ -16,7 +16,7 @@
import dataclasses
import enum
import inspect
from typing import Any, Callable, Protocol, runtime_checkable
from typing import Any, Callable, Protocol, get_type_hints, runtime_checkable
import mujoco
from mujoco.experimental.studio import sim
@@ -204,11 +204,28 @@ def handler(
)
message_type = params[1].annotation
# When `from __future__ import annotations` is enabled or string forward
# references are used, annotations are stored as strings rather than code
# objects. In this case, resolve the string into its actual class type.
if isinstance(message_type, str):
try:
hints = get_type_hints(target)
param_name = params[1].name
if param_name in hints:
message_type = hints[param_name]
except Exception: # pylint: disable=broad-except
pass
# Verify we got a type annotation.
if message_type is inspect.Parameter.empty:
raise TypeError(
f'{fn_name}: second parameter must have a type annotation'
)
# Verify that the annotation resolved to an actual Python class and that the
# class is a valid subclass of Message. This also catches cases where a
# string forward reference failed to resolve via get_type_hints above.
if not (
isinstance(message_type, type) and issubclass(message_type, Message)
):