"""函数注册表.
提供全局函数注册机制, 供 YAML 任务编排通过 ``fn`` 字段引用 Python 函数.
使用方式
--------
import pyflowx as px
@px.register_fn("pack_source")
def pack_source(project_dir, output_dir):
...
# YAML 中引用:
# jobs:
# pack:
# fn: pack_source
# args: ["./project", "./dist"]
"""
from __future__ import annotations
import sys
from typing import Any, Callable, TypeVar, overload
if sys.version_info >= (3, 10):
from typing import ParamSpec
else:
from typing_extensions import ParamSpec # pragma: no cover
__all__ = ["FnRegistry", "get_fn", "has_fn", "register_fn"]
P = ParamSpec("P")
T = TypeVar("T")
_REGISTRY: dict[str, Callable[..., Any]] = {}
@overload
def register_fn(name: Callable[P, T]) -> Callable[P, T]: ...
@overload
def register_fn(name: str | None = None) -> Callable[[Callable[P, T]], Callable[P, T]]: ...
[文档]
def register_fn(name: str | Callable[..., Any] | None = None) -> Callable[..., Any]:
"""装饰器：将函数注册到全局 registry.
支持两种用法::
@register_fn # 使用函数 __name__ 作为注册名
def my_func(): ...
@register_fn("custom") # 显式指定注册名
def my_func(): ...
Parameters
----------
name : str | Callable | None
注册名或被装饰函数; 为 None 时使用函数 ``__name__``
Returns
-------
Callable
装饰器函数或被装饰函数
Raises
------
ValueError
名称已注册或无法推断函数名
"""
if callable(name):
fn = name
key = getattr(fn, "__name__", None)
if key is None:
raise ValueError("无法推断函数名, 请显式提供 name 参数")
if key in _REGISTRY:
raise ValueError(f"函数 {key!r} 已注册")
_REGISTRY[key] = fn
return fn
def decorator(fn: Callable[P, T]) -> Callable[P, T]:
key = name if name is not None else getattr(fn, "__name__", None)
if key is None:
raise ValueError("无法推断函数名, 请显式提供 name 参数")
if key in _REGISTRY:
raise ValueError(f"函数 {key!r} 已注册")
_REGISTRY[key] = fn
return fn
return decorator
[文档]
def get_fn(name: str) -> Callable[..., Any]:
"""按名称获取已注册的函数.
Parameters
----------
name : str
函数名
Returns
-------
Callable
已注册的函数
Raises
------
KeyError
函数未注册
"""
if name not in _REGISTRY:
raise KeyError(f"函数 {name!r} 未注册")
return _REGISTRY[name]
[文档]
def has_fn(name: str) -> bool:
"""检查函数是否已注册.
Parameters
----------
name : str
函数名
Returns
-------
bool
是否已注册
"""
return name in _REGISTRY
class FnRegistry:
"""函数注册表的面向对象访问接口."""
@staticmethod
def register(name: str | None = None) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""注册装饰器, 等价于 :func:`register_fn`."""
return register_fn(name)
@staticmethod
def get(name: str) -> Callable[..., Any]:
"""获取已注册函数, 等价于 :func:`get_fn`."""
return get_fn(name)
@staticmethod
def has(name: str) -> bool:
"""检查是否已注册, 等价于 :func:`has_fn`."""
return has_fn(name)
@staticmethod
def clear() -> None:
"""清空注册表."""
_REGISTRY.clear()
@staticmethod
def names() -> list[str]:
"""返回所有已注册函数名."""
return list(_REGISTRY.keys())