open_pub/Uni-Lab-OS
1
0
Fork 0
mirror of https://github.com/deepmodeling/Uni-Lab-OS synced 2026-04-27 16:35:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support display_name & desc in new registry system

Browse Source 
(cherry picked from commit f71ea2a258)
This commit is contained in:
Xuwznln
2026-04-27 20:28:18 +08:00
parent 195fad9398
commit f145dc04bb
5 changed files with 583 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
unilabos/registry/ast_registry_scanner.py
View File

@@ -32,7 +32,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
MAX_SCAN_DEPTH = 10 # 最大目录递归深度
MAX_SCAN_FILES = 1000 # 最大扫描文件数量
_CACHE_VERSION = 1 # 缓存格式版本号，格式变更时递增
_CACHE_VERSION = 2 # 缓存格式版本号，格式变更时递增
# 合法的装饰器来源模块
_REGISTRY_DECORATOR_MODULE = "unilabos.registry.decorators"
@@ -258,8 +258,6 @@ def scan_directory(
}
# ---------------------------------------------------------------------------
# File-level parsing
# ---------------------------------------------------------------------------
@@ -361,6 +359,7 @@ def _parse_file(
"actions": class_body.get("actions", {}),
"status_properties": class_body.get("status_properties", {}),
"init_params": class_body.get("init_params", []),
"init_docstring": class_body.get("init_docstring"),
"auto_methods": class_body.get("auto_methods", {}),
"import_map": import_map,
}
@@ -497,7 +496,6 @@ def _collect_imports(tree: ast.Module, module_path: str = "") -> Dict[str, str]:
return import_map
# ---------------------------------------------------------------------------
# Decorator finding & argument extraction
# ---------------------------------------------------------------------------
@@ -768,6 +766,7 @@ def _extract_class_body(
"actions": {}, # method_name -> action_info
"status_properties": {}, # prop_name -> status_info
"init_params": [], # [{"name": ..., "type": ..., "default": ...}, ...]
"init_docstring": None,
"auto_methods": {}, # method_name -> method_info (no @action decorator)
}
@@ -780,6 +779,7 @@ def _extract_class_body(
# --- __init__ ---
if method_name == "__init__":
result["init_params"] = _extract_method_params(item, import_map)
result["init_docstring"] = ast.get_docstring(item)
continue
# --- Skip private/dunder ---

120
unilabos/registry/registry.py
View File

@@ -271,6 +271,7 @@ class Registry:
registry_cache.pkl 一个文件中，删除即可完全重置。
"""
import time as _time
from unilabos.registry.ast_registry_scanner import _CACHE_VERSION as AST_SCAN_CACHE_VERSION
from unilabos.registry.ast_registry_scanner import scan_directory
scan_t0 = _time.perf_counter()
@@ -286,6 +287,10 @@ class Registry:
# ---- 统一缓存：一个 pkl 包含所有数据 ----
unified_cache = self._load_config_cache()
ast_cache = unified_cache.setdefault("_ast_scan", {"files": {}})
if ast_cache.get("version") != AST_SCAN_CACHE_VERSION:
ast_cache = {"version": AST_SCAN_CACHE_VERSION, "files": {}}
unified_cache["_ast_scan"] = ast_cache
unified_cache.pop("_build_results", None)
# 默认：扫描 unilabos 包所在的父目录
pkg_root = Path(__file__).resolve().parent.parent # .../unilabos
@@ -561,13 +566,38 @@ class Registry:
return prop_schema
@staticmethod
def _apply_docstring_param_metadata(
schema: Dict[str, Any],
doc_info: Dict[str, Any],
field_to_param: Optional[Dict[str, str]] = None,
) -> None:
"""Apply parsed docstring display names and descriptions to schema properties."""
if not schema or not doc_info:
return
props = schema.get("properties", {})
if not isinstance(props, dict):
return
param_descs = doc_info.get("params", {}) or {}
param_display_names = doc_info.get("param_display_names", {}) or {}
for field_name, prop_schema in props.items():
if not isinstance(prop_schema, dict):
continue
param_name = field_to_param.get(field_name, field_name) if field_to_param else field_name
if not isinstance(param_name, str):
continue
param_name = param_name.removesuffix("[]")
prop_schema["title"] = param_display_names.get(param_name, prop_schema.get("title") or field_name)
prop_schema["description"] = param_descs.get(param_name, prop_schema.get("description") or "")
def _generate_unilab_json_command_schema(
self, method_args: list, docstring: Optional[str] = None,
import_map: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""根据方法参数和 docstring 生成 UniLabJsonCommand schema"""
doc_info = parse_docstring(docstring)
param_descs = doc_info.get("params", {})
schema = {
"type": "object",
@@ -598,12 +628,10 @@ class Registry:
param_name, param_type, param_default, import_map=import_map
)
if param_name in param_descs:
schema["properties"][param_name]["description"] = param_descs[param_name]
if param_required:
schema["required"].append(param_name)
self._apply_docstring_param_metadata(schema, doc_info)
return schema
def _generate_status_types_schema(self, status_methods: Dict[str, Any]) -> Dict[str, Any]:
@@ -799,6 +827,7 @@ class Registry:
type_str = "UniLabJsonCommandAsync" if is_async else "UniLabJsonCommand"
params = method_info.get("params", [])
method_doc = method_info.get("docstring")
method_doc_info = parse_docstring(method_doc)
goal_schema = self._generate_schema_from_ast_params(params, method_name, method_doc, imap)
if action_args is not None:
@@ -828,7 +857,11 @@ class Registry:
# action handles: 从 @action(handles=[...]) 提取并转换为标准格式
raw_handles = (action_args or {}).get("handles")
handles = normalize_ast_action_handles(raw_handles) if isinstance(raw_handles, list) else (raw_handles or {})
handles = (
normalize_ast_action_handles(raw_handles)
if isinstance(raw_handles, list)
else (raw_handles or {})
)
# placeholder_keys: 先从参数类型自动检测，再用装饰器显式配置覆盖/补充
pk = detect_placeholder_keys(params)
@@ -847,7 +880,12 @@ class Registry:
"goal": goal,
"feedback": (action_args or {}).get("feedback") or {},
"result": (action_args or {}).get("result") or {},
"schema": wrap_action_schema(goal_schema, action_name, result_schema=result_schema),
"schema": wrap_action_schema(
goal_schema,
action_name,
description=(action_args or {}).get("description") or method_doc_info.get("description", ""),
result_schema=result_schema,
),
"goal_default": goal_default,
"handles": handles,
"placeholder_keys": pk,
@@ -886,7 +924,11 @@ class Registry:
action_name = f"auto-{action_name}"
raw_handles = action_args.get("handles")
handles = normalize_ast_action_handles(raw_handles) if isinstance(raw_handles, list) else (raw_handles or {})
handles = (
normalize_ast_action_handles(raw_handles)
if isinstance(raw_handles, list)
else (raw_handles or {})
)
method_params = method_info.get("params", [])
@@ -979,7 +1021,10 @@ class Registry:
"schema": schema,
"goal_default": goal_default,
"handles": handles,
"placeholder_keys": {**detect_placeholder_keys(method_params), **(action_args.get("placeholder_keys") or {})},
"placeholder_keys": {
**detect_placeholder_keys(method_params),
**(action_args.get("placeholder_keys") or {}),
},
}
if action_args.get("always_free") or method_info.get("always_free"):
action_entry["always_free"] = True
@@ -988,13 +1033,21 @@ class Registry:
nt = normalize_enum_value(action_args.get("node_type"), NodeType)
if nt:
action_entry["node_type"] = nt
goal_schema_for_docs = action_entry.get("schema", {}).get("properties", {}).get("goal", {})
self._apply_docstring_param_metadata(
goal_schema_for_docs,
parse_docstring(method_info.get("docstring")),
goal,
)
action_value_mappings[action_name] = action_entry
action_value_mappings = dict(sorted(action_value_mappings.items()))
# --- init_param_schema = { config: <init_params>, data: <status_types> } ---
init_params = ast_meta.get("init_params", [])
config_schema = self._generate_schema_from_ast_params(init_params, "__init__", import_map=imap)
config_schema = self._generate_schema_from_ast_params(
init_params, "__init__", ast_meta.get("init_docstring"), import_map=imap
)
data_schema = self._generate_status_schema_from_ast(
ast_meta.get("status_properties", {}), imap
)
@@ -1042,7 +1095,6 @@ class Registry:
) -> Dict[str, Any]:
"""Generate JSON Schema from AST-extracted parameter list."""
doc_info = parse_docstring(docstring)
param_descs = doc_info.get("params", {})
schema: Dict[str, Any] = {
"type": "object",
@@ -1072,12 +1124,10 @@ class Registry:
pname, ptype, pdefault, import_map
)
if pname in param_descs:
schema["properties"][pname]["description"] = param_descs[pname]
if prequired:
schema["required"].append(pname)
self._apply_docstring_param_metadata(schema, doc_info)
return schema
def _generate_status_schema_from_ast(
@@ -1807,7 +1857,7 @@ class Registry:
else:
action_key = f"auto-{k}"
goal_schema = self._generate_unilab_json_command_schema(
v["args"], import_map=enhanced_import_map
v["args"], docstring=v.get("docstring"), import_map=enhanced_import_map
)
ret_type = v.get("return_type", "")
result_schema = None
@@ -1816,7 +1866,13 @@ class Registry:
"result", ret_type, None, import_map=enhanced_import_map
)
old_cfg = old_action_configs.get(action_key) or old_action_configs.get(f"auto-{k}", {})
new_schema = wrap_action_schema(goal_schema, action_key, result_schema=result_schema)
doc_info = parse_docstring(v.get("docstring"))
new_schema = wrap_action_schema(
goal_schema,
action_key,
description=doc_info.get("description", ""),
result_schema=result_schema,
)
old_schema = old_cfg.get("schema", {})
if old_schema:
preserve_field_descriptions(new_schema, old_schema)
@@ -1882,6 +1938,12 @@ class Registry:
merged_pk = dict(old_cfg.get("placeholder_keys", {}))
merged_pk.update(detect_placeholder_keys(v["args"]))
goal_schema_for_docs = (
entry_schema.get("properties", {}).get("goal", {})
if isinstance(entry_schema, dict)
else {}
)
self._apply_docstring_param_metadata(goal_schema_for_docs, doc_info, entry_goal)
entry = {
"type": entry_type,
@@ -1902,7 +1964,8 @@ class Registry:
device_config["init_param_schema"] = {}
init_schema = self._generate_unilab_json_command_schema(
enhanced_info["init_params"], "__init__",
enhanced_info["init_params"],
docstring=enhanced_info.get("init_docstring"),
import_map=enhanced_import_map,
)
device_config["init_param_schema"]["config"] = init_schema
@@ -1949,7 +2012,9 @@ class Registry:
action_str_type_mapping[action_type_str] = target_type
if target_type is not None:
try:
action_config["goal_default"] = ROS2MessageInstance(target_type.Goal()).get_python_dict()
action_config["goal_default"] = ROS2MessageInstance(
target_type.Goal()
).get_python_dict()
except Exception:
action_config["goal_default"] = {}
prev_schema = action_config.get("schema", {})
@@ -2141,10 +2206,15 @@ class Registry:
"unilabos_device_id": {
"type": "string",
"default": "",
"title": "设备ID",
"description": "UniLabOS设备ID，用于指定执行动作的具体设备实例",
},
**schema["properties"]["goal"]["properties"],
}
for field_name, field_schema in schema["properties"]["goal"]["properties"].items():
if isinstance(field_schema, dict):
field_schema.setdefault("title", field_name)
field_schema.setdefault("description", "")
# 将 placeholder_keys 信息添加到 schema 中
if "placeholder_keys" in action_config and action_config.get("schema", {}).get(
"properties", {}
@@ -2212,7 +2282,14 @@ class Registry:
lab_registry = Registry()
def build_registry(registry_paths=None, devices_dirs=None, upload_registry=False, check_mode=False, complete_registry=False, external_only=False):
def build_registry(
registry_paths=None,
devices_dirs=None,
upload_registry=False,
check_mode=False,
complete_registry=False,
external_only=False,
):
"""
构建或获取Registry单例实例
"""
@@ -2226,7 +2303,12 @@ def build_registry(registry_paths=None, devices_dirs=None, upload_registry=False
if path not in current_paths:
lab_registry.registry_paths.append(path)
lab_registry.setup(devices_dirs=devices_dirs, upload_registry=upload_registry, complete_registry=complete_registry, external_only=external_only)
lab_registry.setup(
devices_dirs=devices_dirs,
upload_registry=upload_registry,
complete_registry=complete_registry,
external_only=external_only,
)
# 将 AST 扫描的字符串类型替换为实际 ROS2 消息类（仅查找 ROS2 类型，不 import 设备模块）
lab_registry.resolve_all_types()

81
unilabos/registry/utils.py
View File

@@ -36,16 +36,40 @@ class ROSMsgNotFound(Exception):
# ---------------------------------------------------------------------------
_SECTION_RE = re.compile(r"^(\w[\w\s]*):\s*$")
_PARAM_HEADER_RE = re.compile(
r"^\s*(?P<name>\w[\w]*)\s*(?:\[(?P<display_name>[^\]]+)\])?(?:\s*\([^)]*\))?\s*$"
)
def _parse_docstring_param_header(param_part: str) -> Tuple[str, Optional[str]]:
"""Parse ``name[display_name]`` or Google-style ``name (type)``."""
match = _PARAM_HEADER_RE.match(param_part.strip())
if not match:
return param_part.strip().split("(")[0].strip(), None
display_name = match.group("display_name")
if display_name is not None:
display_name = display_name.strip() or None
return match.group("name").strip(), display_name
def parse_docstring(docstring: Optional[str]) -> Dict[str, Any]:
"""
解析 Google-style docstring，提取描述和参数说明。
解析 docstring，提取描述和参数说明。
支持:
- Google-style ``Args:`` / ``Parameters:`` 小节
- 直接参数行 ``field: desc``
- 带显示名参数行 ``field[Display Name]: desc``
Returns:
{"description": "短描述", "params": {"param1": "参数1描述", ...}}
{
"description": "短描述",
"params": {"param1": "参数1描述", ...},
"param_display_names": {"param1": "显示名", ...},
}
"""
result: Dict[str, Any] = {"description": "", "params": {}}
result: Dict[str, Any] = {"description": "", "params": {}, "param_display_names": {}}
if not docstring:
return result
@@ -53,33 +77,53 @@ def parse_docstring(docstring: Optional[str]) -> Dict[str, Any]:
if not lines:
return result
result["description"] = lines[0].strip()
in_args = False
current_section: Optional[str] = None
current_param: Optional[str] = None
current_display_name: Optional[str] = None
current_desc_parts: list = []
for line in lines[1:]:
def flush_current_param() -> None:
nonlocal current_param, current_display_name, current_desc_parts
if current_param is None:
return
result["params"][current_param] = "\n".join(current_desc_parts).strip()
if current_display_name:
result["param_display_names"][current_param] = current_display_name
current_param = None
current_display_name = None
current_desc_parts = []
first_line = lines[0].strip()
start_index = 0
if not _SECTION_RE.match(first_line) and ":" not in first_line:
result["description"] = first_line
start_index = 1
for line in lines[start_index:]:
stripped = line.strip()
if not stripped:
if current_param is not None:
current_desc_parts.append("")
continue
section_match = _SECTION_RE.match(stripped)
if section_match:
if current_param is not None:
result["params"][current_param] = "\n".join(current_desc_parts).strip()
current_param = None
current_desc_parts = []
section_name = section_match.group(1).lower()
in_args = section_name in ("args", "arguments", "parameters", "params")
flush_current_param()
current_section = section_match.group(1).lower()
in_args = current_section in ("args", "arguments", "parameters", "params")
continue
if not in_args:
parse_as_param = in_args or current_section is None
if not parse_as_param:
continue
if ":" in stripped and not stripped.startswith(" "):
if current_param is not None:
result["params"][current_param] = "\n".join(current_desc_parts).strip()
if ":" in stripped:
flush_current_param()
param_part, _, desc_part = stripped.partition(":")
param_name = param_part.strip().split("(")[0].strip()
param_name, display_name = _parse_docstring_param_header(param_part)
current_param = param_name
current_display_name = display_name
current_desc_parts = [desc_part.strip()]
elif current_param is not None:
aline = line
@@ -89,8 +133,7 @@ def parse_docstring(docstring: Optional[str]) -> Dict[str, Any]:
aline = aline[1:]
current_desc_parts.append(aline.strip())
if current_param is not None:
result["params"][current_param] = "\n".join(current_desc_parts).strip()
flush_current_param()
return result