diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/check_stubs.py | 172 |
1 files changed, 172 insertions, 0 deletions
diff --git a/scripts/check_stubs.py b/scripts/check_stubs.py new file mode 100644 index 0000000..5d5d7e4 --- /dev/null +++ b/scripts/check_stubs.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Check that stubs/ match the real Open WebUI signatures. + +Fetches the relevant Open WebUI source files from GitHub at a given ref +(default: the tag from `required_open_webui_version` in comfier.py), parses +them with `ast` (nothing is imported or installed) and compares function +signatures with the .pyi stubs: async-ness, parameter names/kinds and whether +each parameter has a default. Annotations are not compared. +""" + +import argparse +import ast +import pathlib +import re +import sys +import urllib.request + +ROOT = pathlib.Path(__file__).parent.parent +RAW = "https://raw.githubusercontent.com/open-webui/open-webui/{ref}/backend/open_webui/{path}.py" + +# (stub file, upstream module path, owner class or None, function names) +CHECKS: list[tuple[str, str, str | None, list[str]]] = [ + ("models/config", "models/config", "Config", ["get"]), + ("models/users", "models/users", "UsersTable", ["get_user_by_id"]), + ("routers/images", "routers/images", None, ["get_image_config", "upload_image"]), +] + +Fn = ast.FunctionDef | ast.AsyncFunctionDef + + +def default_ref() -> str: + header = (ROOT / "comfier.py").read_text(encoding="utf-8") + match = re.search(r"^required_open_webui_version:\s*(\S+)", header, re.MULTILINE) + if not match: + sys.exit("Could not find required_open_webui_version in comfier.py") + return "v" + match.group(1).lstrip("v") + + +def fetch(ref: str, path: str) -> str: + url = RAW.format(ref=ref, path=path) + with urllib.request.urlopen(url) as resp: + return resp.read().decode("utf-8") + + +def find_fn(tree: ast.Module, owner: str | None, name: str) -> Fn | None: + body: list[ast.stmt] = tree.body + if owner is not None: + cls = next( + (n for n in body if isinstance(n, ast.ClassDef) and n.name == owner), None + ) + if cls is None: + return None + body = cls.body + return next( + ( + n + for n in body + if isinstance(n, ast.FunctionDef | ast.AsyncFunctionDef) and n.name == name + ), + None, + ) + + +def signature(fn: Fn, drop_self: bool) -> tuple[bool, list[tuple[str, str, bool]]]: + """(is_async, [(name, kind, has_default)])""" + a = fn.args + params: list[tuple[str, str, bool]] = [] + positional = [*a.posonlyargs, *a.args] + n_defaults = len(a.defaults) + for i, arg in enumerate(positional): + has_default = i >= len(positional) - n_defaults + kind = "posonly" if arg in a.posonlyargs else "pos" + params.append((arg.arg, kind, has_default)) + if a.vararg: + params.append((a.vararg.arg, "vararg", False)) + for arg, default in zip(a.kwonlyargs, a.kw_defaults, strict=True): + params.append((arg.arg, "kwonly", default is not None)) + if a.kwarg: + params.append((a.kwarg.arg, "kwarg", False)) + if drop_self and params and params[0][0] in ("self", "cls"): + params = params[1:] + return isinstance(fn, ast.AsyncFunctionDef), params + + +def check_users_singleton(upstream: ast.Module, stub: ast.Module) -> list[str]: + """`Users` must be an instance of UsersTable in both.""" + problems: list[str] = [] + real = any( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "Users" for t in n.targets) + and isinstance(n.value, ast.Call) + and isinstance(n.value.func, ast.Name) + and n.value.func.id == "UsersTable" + for n in upstream.body + ) + stubbed = any( + isinstance(n, ast.AnnAssign) + and isinstance(n.target, ast.Name) + and n.target.id == "Users" + and isinstance(n.annotation, ast.Name) + and n.annotation.id == "UsersTable" + for n in stub.body + ) + if not real: + problems.append("models/users: upstream no longer has `Users = UsersTable()`") + if not stubbed: + problems.append("models/users: stub lacks `Users: UsersTable`") + return problems + + +def is_static(fn: Fn) -> bool: + return any( + isinstance(d, ast.Name) and d.id == "staticmethod" for d in fn.decorator_list + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=(__doc__ or "").splitlines()[0]) + parser.add_argument( + "--ref", default=None, help="git ref/tag (default: from comfier.py)" + ) + args = parser.parse_args() + ref = args.ref or default_ref() + print(f"Checking stubs against open-webui @ {ref}") + + problems: list[str] = [] + for stub_path, upstream_path, owner, names in CHECKS: + stub_file = ROOT / "stubs" / "open_webui" / f"{stub_path}.pyi" + stub_tree = ast.parse(stub_file.read_text(encoding="utf-8")) + try: + upstream_tree = ast.parse(fetch(ref, upstream_path)) + except Exception as err: + problems.append(f"{upstream_path}: could not fetch/parse upstream ({err})") + continue + + for name in names: + label = f"{upstream_path}:{owner + '.' if owner else ''}{name}" + real_fn = find_fn(upstream_tree, owner, name) + stub_fn = find_fn(stub_tree, owner, name) + if real_fn is None: + problems.append(f"{label}: missing upstream") + continue + if stub_fn is None: + problems.append(f"{label}: missing in stub") + continue + real_sig = signature( + real_fn, drop_self=owner is not None and not is_static(real_fn) + ) + stub_sig = signature( + stub_fn, drop_self=owner is not None and not is_static(stub_fn) + ) + if real_sig != stub_sig: + problems.append( + f"{label}:\n upstream: {real_sig}\n stub: {stub_sig}" + ) + else: + print(f" ok {label}") + + if owner == "UsersTable": + problems += check_users_singleton(upstream_tree, stub_tree) + + if problems: + print("\nMismatches:") + for p in problems: + print(f" - {p}") + return 1 + print("All stubs match.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) |