diff --git a/README.md b/README.md index aeaba2f..9c99a95 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,9 @@ gpu.state.blend_set('ALPHA') - **Literal enum types** — string parameters like `gpu.state.blend_set(mode)` use `Literal["NONE", "ALPHA", ...]` instead of plain `str` - **Docstrings** — inline documentation on properties, methods, and functions - **Classmethod detection** — `Matrix.Identity()`, `Vector.Fill()`, etc. correctly typed as `@classmethod` +- **Dunder methods** — `__getitem__`, `__len__`, `__add__`, `__matmul__`, and other operators introspected with correct return types (e.g. `Matrix[0]` returns `Vector`) - **Operator metadata** — `bpy.ops.mesh.primitive_cube_add.poll()`, `.idname()`, `.get_rna_type()`, and `.bl_options` are typed via introspected `_BPyOpsSubModOp` wrapper +- **GPU uniform types** — `shader.uniform_float()` accepts `Matrix`, `Vector`, `Euler`, `Quaternion`, and `Color` in addition to `float | Sequence[float]`, matching the C implementation - **Dynamic array types** — `Image.pixels` typed as `bpy_prop_array[float]` (not `list[float]` or `float`) - **Zero `Any` usage** — precise types throughout, no `typing.Any` fallbacks - **basedpyright strict mode** — 0 errors on generated stubs (4.1+) diff --git a/generate_stubs.py b/generate_stubs.py index 1fce837..05465aa 100644 --- a/generate_stubs.py +++ b/generate_stubs.py @@ -760,18 +760,38 @@ def load_overrides(overrides_dir: str, module_name: str) -> dict[str, ParamOverr return {} +def _apply_func_overrides( + func: FunctionData, func_overrides: ParamOverrides +) -> None: + """Apply overrides to a single function/method.""" + param_overrides = func_overrides.get("params", {}) + for param in func["params"]: + if param["name"] in param_overrides: + param["type"] = param_overrides[param["name"]] + if "return_type" in func_overrides: + func["return_type"] = func_overrides["return_type"] + + def apply_overrides( module_data: ModuleData, overrides: dict[str, ParamOverrides] ) -> ModuleData: - """Apply type overrides to introspected module data.""" + """Apply type overrides to introspected module data. + + Keys can be function names (e.g. ``"my_func"``) for module-level functions, + or ``"ClassName.method_name"`` for struct methods. + """ for func in module_data["functions"]: func_overrides = overrides.get(func["name"], {}) - param_overrides = func_overrides.get("params", {}) - for param in func["params"]: - if param["name"] in param_overrides: - param["type"] = param_overrides[param["name"]] - if "return_type" in func_overrides: - func["return_type"] = func_overrides["return_type"] + if func_overrides: + _apply_func_overrides(func, func_overrides) + + for struct in module_data.get("structs", []): + for method in struct["methods"]: + key = f"{struct['name']}.{method['name']}" + method_overrides = overrides.get(key, {}) + if method_overrides: + _apply_func_overrides(method, method_overrides) + return module_data diff --git a/introspect.py b/introspect.py index 54d0dbb..dcc7f87 100644 --- a/introspect.py +++ b/introspect.py @@ -226,7 +226,7 @@ def parse_docstring_types(docstring: str) -> tuple[dict[str, str], str | None]: # Match :type param: ... up to the next RST directive (:arg, :type, :rtype, :return) # but NOT :class: or :func: which appear inside type annotations - directive_lookahead = r"(?=\n\s*:(?:arg|param|type|rtype|return|returns)[\s:]|$)" + directive_lookahead = r"(?=\n\s*:(?:arg|param|type|rtype|return|returns|raises)[\s:]|$)" for match in re.finditer( rf":type\s+(\w+):\s*(.+?){directive_lookahead}", docstring, re.DOTALL ): @@ -240,6 +240,15 @@ def parse_docstring_types(docstring: str) -> tuple[dict[str, str], str | None]: if rtype_match: return_type = clean_type_str(rtype_match.group(1).strip()) + # Also match standalone `:type: X` (without a param name) used in + # property docstrings (e.g. `:type: bool`, `:type: :class:`Vector``). + if return_type is None: + bare_type_match = re.search( + rf"(? str: """Clean up RST type annotations to plain Python type strings.""" + # Strip "(readonly)" / "(never None)" annotations from Blender docstrings + type_str = re.sub(r",?\s*\(readonly\)", "", type_str) + type_str = re.sub(r",?\s*\(never None\)", "", type_str) type_str = re.sub(r":class:`([^`]+)`", r"\1", type_str) # Remove double backtick RST markup type_str = re.sub(r"``([^`]+)``", r"\1", type_str) @@ -1094,6 +1106,191 @@ def _parse_class_constructor(class_doc: str, cls: type) -> FunctionData | None: } +# Dunders worth exposing in stubs — these affect how the type is used +# in type checking (subscript, iteration, arithmetic, comparison, etc.) +_USEFUL_DUNDERS = { + "__getitem__", + "__setitem__", + "__delitem__", + "__len__", + "__iter__", + "__contains__", + "__add__", + "__radd__", + "__iadd__", + "__sub__", + "__rsub__", + "__isub__", + "__mul__", + "__rmul__", + "__imul__", + "__matmul__", + "__rmatmul__", + "__imatmul__", + "__truediv__", + "__rtruediv__", + "__itruediv__", + "__neg__", + "__pos__", + "__invert__", + "__eq__", + "__ne__", + "__lt__", + "__le__", + "__gt__", + "__ge__", +} + + +def _type_name(obj: object) -> str: + """Return the type name of an object, normalizing NoneType to None.""" + name = type(obj).__name__ + return "None" if name == "NoneType" else name + + +def _fix_dunder_signatures_with_instance( + instance: object, methods: list[FunctionData] +) -> None: + """Fix dunder signatures using an existing instance (for non-constructible types).""" + _fix_dunder_signatures(type(instance), methods, instance=instance) + + +def _fix_dunder_signatures( + cls: type[object], + methods: list[FunctionData], + instance: object | None = None, +) -> None: + """Fix dunder method signatures by runtime probing. + + C wrapper descriptors (``__getitem__``, ``__add__``, etc.) have no type + info in their docstrings. We create a default instance of the class and + call the dunders to discover actual return types and refine parameter types. + """ + + # Fixed return types that don't need runtime probing + _FIXED_RETURNS: dict[str, str] = { + "__len__": "int", + "__contains__": "bool", + "__eq__": "bool", + "__ne__": "bool", + "__lt__": "bool", + "__le__": "bool", + "__gt__": "bool", + "__ge__": "bool", + "__delitem__": "None", + "__setitem__": "None", + } + for method in methods: + fixed = _FIXED_RETURNS.get(method["name"]) + if fixed is not None: + method["return_type"] = fixed + if method["name"] == "__len__": + method["params"] = [] + if method["name"] == "__delitem__": + method["params"] = [ + { + "name": "key", + "type": "int", + "default": None, + "kind": "POSITIONAL_OR_KEYWORD", + } + ] + + if instance is None: + try: + instance = cls() + except Exception: + return + + for method in methods: + name = method["name"] + + # __getitem__: probe with int index to discover element type + if name == "__getitem__": + try: + result = instance[0] + rtype = _type_name(result) + method["return_type"] = rtype + method["params"] = [ + { + "name": "key", + "type": "int", + "default": None, + "kind": "POSITIONAL_OR_KEYWORD", + } + ] + except Exception: + pass + continue + + # __setitem__: refine value type from __getitem__ return type + if name == "__setitem__": + try: + result = instance[0] + vtype = _type_name(result) + method["params"] = [ + { + "name": "key", + "type": "int", + "default": None, + "kind": "POSITIONAL_OR_KEYWORD", + }, + { + "name": "value", + "type": vtype, + "default": None, + "kind": "POSITIONAL_OR_KEYWORD", + }, + ] + except Exception: + pass + continue + + # __neg__, __pos__, __invert__: unary ops return same type + if name in ("__neg__", "__pos__", "__invert__"): + try: + op = getattr(instance, name) + result = op() + method["return_type"] = _type_name(result) + method["params"] = [] + except Exception: + pass + continue + + # Binary arithmetic: probe with same-type operand to get return type. + # The parameter type stays as object (could be Self, float, etc.) + if name in ( + "__add__", + "__radd__", + "__iadd__", + "__sub__", + "__rsub__", + "__isub__", + "__mul__", + "__rmul__", + "__imul__", + "__matmul__", + "__rmatmul__", + "__imatmul__", + "__truediv__", + "__rtruediv__", + "__itruediv__", + ): + try: + op = getattr(instance, name) + result = op(instance) + method["return_type"] = _type_name(result) + except Exception: + # If same-type fails (e.g. Vector / Vector), try with float + try: + op = getattr(instance, name) + result = op(1.0) + method["return_type"] = _type_name(result) + except Exception: + pass + continue + + def introspect_class(cls: type[object], module_name: str) -> StructData: """Introspect a class (C extension or Python) and return StructData.""" # Determine base class (skip object and internal bases) @@ -1118,7 +1315,7 @@ def introspect_class(cls: type[object], module_name: str) -> StructData: methods: list[FunctionData] = [] for name in sorted(dir(cls)): - if name.startswith("_"): + if name.startswith("_") and name not in _USEFUL_DUNDERS: continue try: @@ -1177,6 +1374,13 @@ def introspect_class(cls: type[object], module_name: str) -> StructData: if init_method: methods.insert(0, init_method) + # Fix dunder signatures by runtime probing. C wrapper descriptors have no + # type info in docstrings, but we can create a default instance and call + # the dunders to discover actual return types and parameter types. + dunder_methods = [m for m in methods if m["name"] in _USEFUL_DUNDERS] + if dunder_methods: + _fix_dunder_signatures(cls, dunder_methods) + return { "name": cls.__name__, "doc": class_doc, @@ -1256,7 +1460,7 @@ def introspect_ops_module() -> ModuleData: continue try: result = func() - _return_type_fixups[name] = type(result).__name__ + _return_type_fixups[name] = _type_name(result) except Exception: pass @@ -1439,6 +1643,53 @@ def introspect_module(module_name: str) -> ModuleData: infer_getter_return_types(functions) + # Discover hidden C types reachable through properties of known classes. + # E.g. mathutils.Matrix.col returns MatrixAccess which isn't in dir(mathutils). + # Only enabled for modules known to have safe, side-effect-free constructors. + _SAFE_PROBE_MODULES = {"mathutils"} + if module_name in _SAFE_PROBE_MODULES: + import builtins as _builtins_mod + + _SKIP_PROP_TYPES = frozenset( + ("int", "float", "str", "bool", "NoneType", "list", "tuple", "dict", "set") + ) + known_struct_names = {s["name"] for s in structs} + for struct in list(structs): + cls = getattr(module, struct["name"], None) + if cls is None or not isinstance(cls, type): + continue + try: + instance = cls() + except Exception: + continue + for prop in struct["properties"]: + if prop["type"] != "object": + continue + try: + val = getattr(instance, prop["name"]) + except Exception: + continue + val_type = type(val) + val_name = val_type.__name__ + if hasattr(_builtins_mod, val_name) or val_name in _SKIP_PROP_TYPES: + continue + # Add the hidden type if not already known + if val_name not in known_struct_names: + hidden_struct = introspect_class(val_type, module_name) + # Fix dunders using the live instance (since the type + # may not be directly constructible) + dunder_methods = [ + m + for m in hidden_struct["methods"] + if m["name"] in _USEFUL_DUNDERS + ] + if dunder_methods: + _fix_dunder_signatures_with_instance(val, dunder_methods) + structs.append(hidden_struct) + known_struct_names.add(val_name) + # Fix the property type either way + prop["type"] = val_name + return { "module": module_name, "doc": inspect.getdoc(module) or "", @@ -1974,7 +2225,12 @@ def introspect_rna_types() -> ModuleData: if cls is not None: struct = introspect_class(cls, "bpy.types") struct["base"] = _GENERIC_BASES[cls_name] - struct["methods"] = _EXTRA_DUNDERS[cls_name] + struct["methods"] + # Replace introspected dunders with manually-typed versions that + # use generic type parameters (_T). + extra_names = {m["name"] for m in _EXTRA_DUNDERS[cls_name]} + struct["methods"] = _EXTRA_DUNDERS[cls_name] + [ + m for m in struct["methods"] if m["name"] not in extra_names + ] structs.append(struct) # Build a map of collection wrapper class -> element type. diff --git a/overrides/5.0/gpu.types.json b/overrides/5.0/gpu.types.json new file mode 100644 index 0000000..cc0037f --- /dev/null +++ b/overrides/5.0/gpu.types.json @@ -0,0 +1,7 @@ +{ + "GPUShader.uniform_float": { + "params": { + "value": "float | int | Sequence[float] | mathutils.Matrix | mathutils.Vector | mathutils.Euler | mathutils.Quaternion | mathutils.Color" + } + } +}