Matija Kecman
f1c8d3a58f
This change refactors the Python API to use the term "plugins" instead of "handlers" for classes containing decorated handler methods. The term handler is still used for an annotated method of a plugin class that handles a specific message type. Also improved some documentation. PiperOrigin-RevId: 962150099 Change-Id: I34a8cc410cd784b088605490cfa3baccea5c9e71
63 lines
2.3 KiB
Python
63 lines
2.3 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Plugin registry, handler discovery, and runtime dispatching."""
|
|
|
|
import collections
|
|
from typing import Any, Callable
|
|
from mujoco.experimental.studio import messages
|
|
|
|
|
|
def _discover_handlers(
|
|
plugin: Any,
|
|
) -> list[tuple[messages._HandlerInfo, Callable[..., Any]]]:
|
|
"""Scan an object for methods marked with handler decorators."""
|
|
handlers = []
|
|
for name in dir(type(plugin)):
|
|
unbound = getattr(type(plugin), name, None)
|
|
if unbound is None:
|
|
continue
|
|
info = getattr(unbound, messages._HANDLER_INFO_ATTR, None) # pylint: disable=protected-access
|
|
if info is None:
|
|
continue
|
|
bound = getattr(plugin, name)
|
|
handlers.append((info, bound))
|
|
return handlers
|
|
|
|
|
|
class PluginRegistry:
|
|
"""Runtime registry of discovered handlers from plugin instances."""
|
|
|
|
def __init__(self, plugins: list[Any] | None = None) -> None:
|
|
self.handlers: dict[
|
|
type[messages.Message], list[tuple[int, Callable[..., Any]]]
|
|
] = collections.defaultdict(list)
|
|
for obj in plugins or []:
|
|
for info, bound_method in _discover_handlers(obj):
|
|
self.handlers[info.message_type].append((info.priority, bound_method))
|
|
|
|
def dispatch(self, message: Any) -> None:
|
|
"""Dispatches a message to registered handlers in priority order."""
|
|
# Collect all matching handlers for the given message class and all its
|
|
# ancestors/superclasses.
|
|
message_handlers = []
|
|
for message_type in type(message).__mro__:
|
|
message_handlers.extend(self.handlers.get(message_type, []))
|
|
|
|
# Sort by priority; higher values execute first.
|
|
message_handlers.sort(key=lambda item: item[0], reverse=True)
|
|
for _, handler in message_handlers:
|
|
if handler(message):
|
|
# Handler returned True to consume the message; stop dispatching.
|
|
return
|