399 lines
14 KiB
Python
399 lines
14 KiB
Python
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)) |