base_agent/uas/uas_control/command_registry.py

399 lines
14 KiB
Python
Raw Normal View History

2026-06-18 03:28:14 +00:00
import functools
import inspect
from dataclasses import dataclass
from typing import Any, Callable, get_type_hints
# ─────────────────────────────────────────────
# 数据结构定义
# ─────────────────────────────────────────────
@dataclass
class ParamInfo:
"""单个参数的元信息"""
name: str
annotation: type | None
default: Any
kind: str
@property
def has_default(self) -> bool:
return self.default is not inspect.Parameter.empty
@property
def type_name(self) -> str:
if self.annotation is None:
return "Any"
return getattr(self.annotation, "__name__", str(self.annotation))
def __repr__(self):
default_str = f" = {self.default!r}" if self.has_default else ""
return f"{self.name}: {self.type_name}{default_str}"
@dataclass
class CommandInfo:
"""函数完整元信息"""
name: str # 注册名
qualname: str # 限定名,如 DroneController.takeoff
module: object
doc: str | None
params: list[ParamInfo]
return_annotation: type | None
func: Callable
is_method: bool = False # 是否为类方法
class_name: str | None = None # 所属类名
method_type: str = "method" # method / classmethod / staticmethod
@property
def return_type_name(self) -> str:
if self.return_annotation is None:
return "Any"
return getattr(self.return_annotation, "__name__", str(self.return_annotation))
def signature(self) -> str:
params_str = ", ".join(repr(p) for p in self.params)
prefix = f"{self.class_name}." if self.class_name else ""
return f"{prefix}{self.name}({params_str}) -> {self.return_type_name}"
def to_dict(self) -> dict:
return {
"name": self.name,
"qualname": self.qualname,
"module": self.module,
"doc": self.doc,
"is_method": self.is_method,
"class_name": self.class_name,
"method_type": self.method_type,
"return_type": self.return_type_name,
"params": [
{
"name": p.name,
"type": p.type_name,
"default": p.default if p.has_default else "<required>",
"kind": p.kind,
}
for p in self.params
],
}
# ─────────────────────────────────────────────
# 注册表(单例)
# ─────────────────────────────────────────────
class CommandRegistry:
"""全局函数注册表"""
_instance: "CommandRegistry | None" = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._registry: dict[str, CommandInfo] = {}
return cls._instance
def register(self, info: CommandInfo) -> None:
self._registry[info.name] = info
def get(self, name: str) -> CommandInfo | None:
return self._registry.get(name)
def all(self) -> dict[str, CommandInfo]:
return dict(self._registry)
def names(self) -> list[str]:
return list(self._registry.keys())
def by_class(self, class_name: str) -> list[CommandInfo]:
"""查询某个类下注册的所有方法"""
return [i for i in self._registry.values() if i.class_name == class_name]
def __repr__(self):
lines = [f"CommandRegistry ({len(self._registry)} entries):"]
# 按 class_name 分组展示
groups: dict[str | None, list[CommandInfo]] = {}
for info in self._registry.values():
groups.setdefault(info.class_name, []).append(info)
for cls_name, infos in groups.items():
label = f"[class {cls_name}]" if cls_name else "[functions]"
lines.append(f" {label}")
for info in infos:
tag = f"@{info.method_type}" if info.method_type != "method" else ""
lines.append(f" {'' if info.is_method else ''} {info.signature()} {tag}")
return "\n".join(lines)
# 全局注册表实例
registry = CommandRegistry()
# ─────────────────────────────────────────────
# 内部解析工具
# ─────────────────────────────────────────────
def _parse_func_info(
func: Callable,
reg_name: str,
skip_params: set[str],
method_type: str,
) -> CommandInfo:
"""从函数对象中提取完整元信息"""
sig = inspect.signature(func)
try:
hints = get_type_hints(func)
except Exception:
hints = {}
params: list[ParamInfo] = []
for param_name, param in sig.parameters.items():
if param_name in skip_params:
continue
params.append(ParamInfo(
name=param_name,
annotation=hints.get(param_name),
default=param.default,
kind=param.kind.name,
))
# 从 qualname 推断类名,如 "MyClass.method" → "MyClass"
parts = func.__qualname__.split(".")
class_name = parts[-2] if len(parts) >= 2 and "<locals>" not in parts[-2] else None
is_method = class_name is not None
return CommandInfo(
name=reg_name,
qualname=func.__qualname__,
module=func,
doc=inspect.getdoc(func),
params=params,
return_annotation=hints.get("return"),
func=func,
is_method=is_method,
class_name=class_name,
method_type=method_type,
)
# ─────────────────────────────────────────────
# 装饰器实现
# ─────────────────────────────────────────────
def register(
name: str | None = None,
*,
override: bool = False,
registry: CommandRegistry = registry
) -> Callable:
"""
函数/方法注册装饰器
支持以下所有场景
- 普通函数
- 实例方法自动跳过 self
- 类方法 @classmethod自动跳过 cls
- 静态方法 @staticmethod
注册名默认为函数名类方法会额外记录所属类名
对于类方法需将 @register 放在 @classmethod / @staticmethod 内侧
Args:
name: 自定义注册名默认使用函数名
override: 是否允许覆盖已注册的同名函数
registry:
Usage:
# 普通函数
@register
def foo(x: int) -> bool: ...
# 实例方法
class MyClass:
@register
def bar(self, x: int) -> None: ...
# 类方法(@register 在内侧)
class MyClass:
@classmethod
@register
def baz(cls, x: int) -> None: ...
# 静态方法(@register 在内侧)
class MyClass:
@staticmethod
@register
def qux(x: int) -> None: ...
# 自定义名称
@register(name="my_func")
def some_func(x: float) -> None: ...
"""
def decorator(func: Callable) -> Callable:
# ── 处理 classmethod / staticmethod 描述符 ──────
# 当 @register 在 @classmethod 内侧时func 是原始函数
# 当 @register 在 @classmethod 外侧时func 是 classmethod 描述符
method_type = "method"
raw_func = func
if isinstance(func, classmethod):
method_type = "classmethod"
raw_func = func.__func__
elif isinstance(func, staticmethod):
method_type = "staticmethod"
raw_func = func.__func__
# 确定注册名
reg_name = name or raw_func.__name__
# 重复注册检查
if not override and reg_name in registry.all():
raise ValueError(
f"[Registry] '{reg_name}' 已注册。"
f"如需覆盖请使用 @register(override=True)"
)
# 确定需要跳过的隐式参数
skip_params: set[str] = set()
if method_type == "method":
skip_params = {"self"}
elif method_type == "classmethod":
skip_params = {"cls"}
# 解析并注册
info = _parse_func_info(raw_func, reg_name, skip_params, method_type)
registry.register(info)
# 将注册信息挂载到原始函数
raw_func.__registry_info__ = info
# ── 根据方法类型返回正确的包装器 ────────────────
if isinstance(func, classmethod):
# 重新包装为 classmethod
@functools.wraps(raw_func)
def cm_wrapper(cls_arg, *args, **kwargs):
return raw_func(cls_arg, *args, **kwargs)
cm_wrapper.__registry_info__ = info
return classmethod(cm_wrapper)
elif isinstance(func, staticmethod):
# 重新包装为 staticmethod
@functools.wraps(raw_func)
def sm_wrapper(*args, **kwargs):
return raw_func(*args, **kwargs)
sm_wrapper.__registry_info__ = info
return staticmethod(sm_wrapper)
else:
# 普通函数 / 实例方法
@functools.wraps(raw_func)
def wrapper(*args, **kwargs):
return raw_func(*args, **kwargs)
wrapper.__registry_info__ = info
return wrapper
# ── 支持 @register 和 @register() 两种写法 ──────────
if callable(name):
func, name = name, None
return decorator(func)
return decorator
# ─────────────────────────────────────────────
# 使用示例
# ─────────────────────────────────────────────
if __name__ == "__main__":
# ── 普通函数 ────────────────────────────────────────
@register
def emergency_stop() -> None:
"""全局紧急停止"""
print("紧急停止!")
# ── 类:实例方法 + 类方法 + 静态方法 ────────────────
class DroneController:
@register
def takeoff(self, altitude: float, speed: float = 2.0) -> bool:
"""控制无人机起飞"""
print(f"起飞至 {altitude}m速度 {speed}m/s")
return True
@register
def goto(self, lat: float, lon: float, alt: float = 50.0) -> None:
"""飞往指定 GPS 坐标"""
print(f"飞往 ({lat}, {lon}),高度 {alt}m")
@register
def land(self) -> bool:
"""降落"""
print("降落中...")
return True
@classmethod
@register # @register 在 @classmethod 内侧
def create(cls, device_id: str) -> "DroneController":
"""工厂方法:创建控制器实例"""
return cls()
@staticmethod
@register # @register 在 @staticmethod 内侧
def ping(host: str, timeout: int = 3) -> bool:
"""检测设备连通性"""
return True
# ── 另一个类 ────────────────────────────────────────
class DogRobot:
@register
def walk(self, speed: float = 1.0, direction: float = 0.0) -> None:
"""四足机器人行走"""
print(f"行走:速度={speed}, 方向={direction}°")
@register(name="dog_trot") # 自定义注册名
def trot(self, speed: float = 2.5) -> None:
"""四足机器人小跑"""
print(f"小跑:速度={speed}")
@staticmethod
@register
def get_supported_gaits() -> list:
"""获取支持的步态列表"""
return ["walk", "trot", "bound"]
# ── 打印注册表 ───────────────────────────────────────
print("=" * 65)
print(registry)
print()
# ── 按类查询 ─────────────────────────────────────────
print("=" * 65)
print("DroneController 注册方法:")
for info in registry.by_class("DroneController"):
print(f" [{info.method_type:11s}] {info.signature()}")
print()
# ── 查询单个方法详情 ──────────────────────────────────
print("=" * 65)
info = registry.get("takeoff")
print(f"函数名: {info.name}")
print(f"限定名: {info.qualname}")
print(f"所属类: {info.class_name}")
print(f"方法类型: {info.method_type}")
print(f"文档: {info.doc}")
print(f"签名: {info.signature()}")
print("参数详情:")
for p in info.params:
print(f" {repr(p)} [kind={p.kind}]")
print()
# ── JSON 导出 ─────────────────────────────────────────
import json
print("=" * 65)
print("JSON 导出takeoff:")
print(json.dumps(info.to_dict(), ensure_ascii=False, indent=2))