From e4284ce7d90466598ced419196efe5ae7a9fa8fd Mon Sep 17 00:00:00 2001 From: Joseph HENRY Date: Fri, 3 Apr 2026 15:54:40 +0200 Subject: [PATCH] Use @property/@setter for writable mathutils properties Writable mathutils properties (Vector, Euler, etc.) were typed as plain attributes with a union type, causing the getter to also return the union. This made `obj.location.x` fail type checking since Sequence[float] has no `.x`. Now these properties use @property for the getter (returning the concrete mathutils type) and @setter accepting the wider union. Co-Authored-By: Claude Opus 4.6 (1M context) --- generate_stubs.py | 36 +++++++++++++++++++++++++++--------- introspect.py | 30 +++++++++++++++++------------- tests/test_generate_stubs.py | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 79 insertions(+), 22 deletions(-) diff --git a/generate_stubs.py b/generate_stubs.py index 524ad55..eaa30ba 100644 --- a/generate_stubs.py +++ b/generate_stubs.py @@ -89,6 +89,9 @@ def collect_all_type_strings(module_data: ModuleData) -> list[str]: all_types.append(base) for prop in struct["properties"]: all_types.append(prop["type"]) + setter_type = prop.get("setter_type") + if setter_type: + all_types.append(setter_type) for method in struct["methods"]: for param in method["params"]: if param["type"]: @@ -367,7 +370,9 @@ def generate_property_stub( property_decorator: str = "@property", ) -> str: """Generate a stub for a class property.""" - if prop["is_readonly"]: + setter_type = prop.get("setter_type") + + if prop["is_readonly"] or setter_type: result = f"{indent}{property_decorator}\n" result += f"{indent}def {prop['name']}(self) -> {prop['type']}:\n" if prop["description"]: @@ -377,6 +382,11 @@ def generate_property_stub( result += f'{indent} """{desc}"""\n' else: result += f"{indent} ...\n" + if setter_type: + result += f"\n{indent}@{prop['name']}.setter\n" + result += ( + f"{indent}def {prop['name']}(self, value: {setter_type}) -> None: ...\n" + ) return result result = f"{indent}{prop['name']}: {prop['type']}\n" @@ -413,6 +423,7 @@ def fixup_shadowed_builtins( fixed: list[PropertyData] = [] for prop in properties: new_type = prop["type"] + new_setter_type = prop.get("setter_type") for name in shadowed: # Replace bare builtin type references with builtins.X # e.g. "int" -> "builtins.int", "list[int]" -> "list[builtins.int]" @@ -421,14 +432,21 @@ def fixup_shadowed_builtins( f"builtins.{name}", new_type, ) - fixed.append( - { - "name": prop["name"], - "type": new_type, - "is_readonly": prop["is_readonly"], - "description": prop["description"], - } - ) + if new_setter_type: + new_setter_type = re.sub( + rf"\b{name}\b", + f"builtins.{name}", + new_setter_type, + ) + new_prop: PropertyData = { + "name": prop["name"], + "type": new_type, + "is_readonly": prop["is_readonly"], + "description": prop["description"], + } + if new_setter_type: + new_prop["setter_type"] = new_setter_type + fixed.append(new_prop) return fixed diff --git a/introspect.py b/introspect.py index 8d4d239..4efcd19 100644 --- a/introspect.py +++ b/introspect.py @@ -194,7 +194,11 @@ class VariableData(TypedDict): value: str -class PropertyData(TypedDict): +class _PropertyDataOptional(TypedDict, total=False): + setter_type: str + + +class PropertyData(_PropertyDataOptional): name: str type: str is_readonly: bool @@ -1552,14 +1556,15 @@ def _property_data_from_member( else: is_readonly = not _is_getset_writable(cls, name) prop_type = rtype or "object" - if not is_readonly and prop_type in _MATHUTILS_ARRAY_TYPES: - prop_type = f"{prop_type} | Sequence[float]" - return { + result: PropertyData = { "name": name, "type": prop_type, "is_readonly": is_readonly, "description": doc, } + if not is_readonly and prop_type in _MATHUTILS_ARRAY_TYPES: + result["setter_type"] = f"{prop_type} | Sequence[float]" + return result def _introspect_declared_class_members( @@ -2661,17 +2666,16 @@ def _rna_struct_to_data( for prop in getattr(struct_info, "properties", []): prop_type = rna_property_to_type(prop) is_readonly = bool(getattr(prop, "is_readonly", False)) + prop_data: PropertyData = { + "name": str(getattr(prop, "identifier", "")), + "type": prop_type, + "is_readonly": is_readonly, + "description": str(getattr(prop, "description", "") or ""), + } # Writable mathutils properties also accept Sequence[float] for assignment if not is_readonly and prop_type in _MATHUTILS_ARRAY_TYPES: - prop_type = f"{prop_type} | Sequence[float]" - properties.append( - { - "name": str(getattr(prop, "identifier", "")), - "type": prop_type, - "is_readonly": is_readonly, - "description": str(getattr(prop, "description", "") or ""), - } - ) + prop_data["setter_type"] = f"{prop_type} | Sequence[float]" + properties.append(prop_data) methods: list[FunctionData] = [] is_collection_wrapper = sid in collection_element_types diff --git a/tests/test_generate_stubs.py b/tests/test_generate_stubs.py index f16f314..6cef3b6 100644 --- a/tests/test_generate_stubs.py +++ b/tests/test_generate_stubs.py @@ -420,6 +420,41 @@ class TestGeneratePropertyStub(unittest.TestCase): self.assertIn("def type(self) -> str:", result) self.assertIn("...", result) + def test_setter_type_generates_property_with_setter(self) -> None: + prop: PropertyData = { + "name": "location", + "type": "mathutils.Vector", + "is_readonly": False, + "description": "Object location.", + "setter_type": "mathutils.Vector | Sequence[float]", + } + result = generate_property_stub(prop) + self.assertIn("@property", result) + self.assertIn("def location(self) -> mathutils.Vector:", result) + self.assertIn('"""Object location."""', result) + self.assertIn("@location.setter", result) + self.assertIn( + "def location(self, value: mathutils.Vector | Sequence[float]) -> None: ...", + result, + ) + + def test_setter_type_no_description(self) -> None: + prop: PropertyData = { + "name": "rotation", + "type": "mathutils.Euler", + "is_readonly": False, + "description": "", + "setter_type": "mathutils.Euler | Sequence[float]", + } + result = generate_property_stub(prop) + self.assertIn("@property", result) + self.assertIn("def rotation(self) -> mathutils.Euler:", result) + self.assertIn("@rotation.setter", result) + self.assertIn( + "def rotation(self, value: mathutils.Euler | Sequence[float]) -> None: ...", + result, + ) + class TestGenerateMethodStub(unittest.TestCase): def test_simple_method(self) -> None: