Stub generation and minor refactoring

.gitignore

@@ -8,4 +8,5 @@
 /src/gdmath/*.cpp
 /src/gdmath/*.pyd
 /src/gdmath/_gdmath.pyx
+/src/gdmath/_gdmath.pyi
 /codegen/output/*

MANIFEST.in

@@ -2,4 +2,5 @@ include src/gdmath/*.c
 include src/gdmath/*.pyd
 include src/gdmath/*.so
 recursive-exclude tests *
-recursive-exclude codegen *
+recursive-exclude codegen *
+prune */__pycache__

README.md

@@ -17,6 +17,7 @@
   - Transform
     - [Transform2D](https://github.com/shBLOCK/GdMath/wiki#transform2d) & [Transform3D](https://github.com/shBLOCK/GdMath/wiki#transform3d)
     - Godot and GLSL like api
+- Stubs for better IDE support
 
 Please refer to the [wiki](https://github.com/shBLOCK/GdMath/wiki) for more details
 

codegen/codegen_helper.py

@@ -352,18 +352,18 @@ def process_overloads(file: str) -> str:
         if is_overload_func:
             is_overload_func = False
             m = regex.search(
-                r"cdef\s*"
+                r"cdef\s+"
                 r"(?P<return>\w+)?\s+"
                 r"(?P<name>\w+)\s*"
                 r"\("
-                r"\s*(?:(?P<first_param>\w+\s+\w+|self)(?:\s*,\s*(?P<remaining_params>\w+\s+\w+))*)?\s*"
+                r"\s*(?:(?P<params>\w+\s+\w+|self)(?:\s*,\s*(?P<params>\w+\s+\w+))*)?\s*"
                 r"\)",
                 line
             )
             assert m is not None, line
             name = m.group("name")
             c_ret = m.group("return")
-            c_params = (*m.captures("first_param"), *m.captures("remaining_params"))
+            c_params = tuple(m.captures("params"))
             c_params = tuple(regex.match(r"\w+", p).group() for p in c_params)
             if name not in overloads:
                 overloads[name] = _Overload(name)
@@ -427,6 +427,16 @@ def step_generate(template_file: str, output_file: str = None, write_file: bool
         return result
 
 
+def step_gen_stub(source_file: str, output_file: str):
+    import stub_generator
+    source = open(f"output/{source_file}", encoding="utf8").read()
+    t = time.perf_counter()
+    result = stub_generator.gen_stub(source)
+    print(f"Step Gen Stub: {source_file} -> {output_file} completed in {time.perf_counter() - t:.3f}s")
+    with open(f"output/{output_file}", "w", encoding="utf8") as output:
+        output.write(result)
+
+
 def step_cythonize(file: str):
     import sys
     import subprocess

codegen/gen_all.py

@@ -4,10 +4,12 @@
 
 if __name__ == '__main__':
     codegen.step_generate("_gdmath.pyx", write_file=True, _globals=globals())
+    codegen.step_gen_stub("_gdmath.pyx", "_gdmath.pyi")
 
     import os
     if os.getenv("CI") != "true":
         codegen.step_move_to_dest("../src/gdmath/", "_gdmath", ".pyx")
+        codegen.step_move_to_dest("../src/gdmath/", "_gdmath", ".pyi")
 
         import sys
         import subprocess

codegen/stub_generator.py

@@ -0,0 +1,241 @@
+from itertools import count
+from typing import Sequence, Optional
+
+import regex
+
+def convert_type(org: str) -> str:
+    types = []
+    for raw in map(str.strip, org.split("|")):
+        match raw:
+            case "float" | "double" | "py_float":
+                types.append("(float | int)")
+            case "int" | "long" | "py_int":
+                types.append("int")
+            case "None" | "void":
+                types.append("None")
+            case "object":
+                types.append("Any")
+            case _:
+                types.append(f'"{raw}"')
+    return types[0] if len(types) == 1 else f"Union[{', '.join(types)}]"
+
+class StubProperty:
+    def __init__(self, name: str, ptype: str, mutable: bool = True):
+        self.name = name
+        self.type = ptype
+        self.mutable = mutable
+
+    def stub(self) -> list[str]:
+        if self.mutable:
+            return [f"{self.name}: {self.type}"]
+        else:
+            return [
+                "@property",
+                f"def {self.name}(self) -> {self.type}: ..."
+            ]
+
+class StubMethod:
+    class Param:
+        def __init__(self, name: str, ptype: str, default: Optional[str] = None, const_mapping: dict[str, str] = None):
+            self.name = name
+            self.ptype = ptype
+            self.default = default
+            self.const_mapping = const_mapping
+
+        def stub(self) -> str:
+            stub = f"{self.name}: {self.ptype}"
+            if self.default is not None:
+                default = self.default
+                if self.const_mapping:
+                    default = self.const_mapping.get(default, default)
+                stub += f" = {default}"
+            return stub
+
+    def __init__(self, name: str, rtype: str, params: Sequence[Param | str], is_cdef: bool, is_static: bool):
+        self.name = name
+        self.rtype = rtype
+        self.params = params
+        self.is_cdef = is_cdef
+        self.is_static = is_static
+
+    def stub(self, name_override: Optional[str] = None) -> list[str]:
+        stub = []
+        if self.is_static:
+            stub.append("@staticmethod")
+
+        param_list = [] if self.is_static else ["self"]
+        for p in self.params:
+            if type(p) is str:
+                param_list.append(p)
+            else:
+                param_list.append(p.stub())
+
+        stub.append(f"def {name_override or self.name}({', '.join(param_list)}) -> {self.rtype}: ...")
+        return stub
+
+class StubClass:
+    def __init__(self, name: str):
+        self.name = name
+        self.properties: dict[str, StubProperty] = {}
+        self.methods: dict[str, StubMethod] = {}
+
+    def stub(self) -> list[str]:
+        def indent(src: list[str]) -> list[str]:
+            for i, line in enumerate(src):
+                src[i] = " " * 4 + line
+            return src
+
+        stub = [
+            "# noinspection SpellCheckingInspection"
+            f"class {self.name}:"
+        ]
+
+        for prop in self.properties.values():
+            stub.extend(indent(prop.stub()))
+
+        stub.append("")
+
+        methods = self.methods.copy()
+        while methods:
+            _, method = methods.popitem()
+            if not method.is_cdef:
+                # Overloaded
+                if f"_{method.name}_0" in methods:
+                    for i in count():
+                        ol_name = f"_{method.name}_{i}"
+                        if ol_method := methods.get(ol_name):
+                            del methods[ol_name]
+                            assert ol_method.is_cdef
+                            stub.extend(indent(["@overload"]))
+                            stub.extend(indent(ol_method.stub(name_override=method.name)))
+                        else:
+                            break
+                # Not Overloaded
+                else:
+                    stub.extend(indent(method.stub()))
+
+        return stub
+
+
+def gen_stub(source: str) -> str:
+    const_mapping = {}
+    decorators = []
+    classes = []
+    clazz = None
+
+    print("gen_stub: reading source...")
+    source_lines = source.splitlines(keepends=False)
+    for line_no, line in enumerate(source_lines):
+        # Class
+        if m := regex.match(r"cdef\s+class\s+(?P<name>\w+)\s*:", line):
+            if clazz is not None:
+                classes.append(clazz)
+            clazz = StubClass(m.group("name"))
+            decorators.clear()
+        # Property
+        elif m := regex.match(r"\s+cdef\s+public\s+(?P<type>\w+)\s+(?P<names>\w+)(?:\s*,\s*(?P<names>\w+))*", line):
+            ptype = convert_type(m.group("type"))
+            for pname in m.captures("names"):
+                clazz.properties[pname] = StubProperty(pname, ptype)
+            decorators.clear()
+        # Method (including @property)
+        elif (
+                (cdef_m := regex.match( # cdef methods can not be static, self always present
+                    r"\s+cdef\s+"
+                    r"(?:inline\s+)?"
+                    r"(?P<return>\w+)?\s+"
+                    r"(?P<name>\w+)\s*"
+                    r"\(\s*"
+                    r"(?:self\s*)?"
+                    r"(?:,\s*(?P<params>\w+\s+\w+)\s*)*"
+                    r"\)\s*"
+                    r"(?:noexcept)?\s*"
+                    r":",
+                    line
+                ))
+                or
+                (def_m := regex.match(
+                    r"\s+def\s+"
+                    r"(?P<name>\w+)\s*"
+                    r"\(\s*"
+                    r"(?:(?:self\s*)|(?&_param))?"
+                    r"(?:,\s*(?P<_param>(?P<params>\w+\s+\w+(?:\s*=\s*[^,)]+)?)|(?P<params>/))\s*)*"
+                    r"\)\s*"
+                    r"(?:->\s*(?P<return>[^:]+))?\s*"
+                    r":",
+                    line
+                ))
+            ):
+            m = cdef_m or def_m
+            is_cdef = cdef_m is not None
+            name = m.group("name")
+
+            assert "classmethod" not in decorators
+            # property getter
+            if "property" in decorators:
+                assert name not in clazz.properties
+                clazz.properties[name] = StubProperty(name, convert_type(m.group("return")), mutable=False)
+            # proeprty setter
+            elif setter_dec := [d for d in decorators if d.endswith(".setter")]:
+                assert len(setter_dec) == 1
+                pname = setter_dec[0].split(".")[0]
+                clazz.properties[pname].mutable = True
+            # normal method
+            else:
+                rtype = convert_type(m.group("return") or "object")
+                if source_lines[line_no + 1].strip() == "#<RETURN_SELF>":
+                    rtype = "Self"
+
+                params = []
+                for param in m.captures("params"):
+                    if param == "/":
+                        params.append(param)
+                        continue
+                    if is_cdef:
+                        # default values are not supported for cdef methods yet
+                        pm = regex.fullmatch(r"(?P<type>\w+)\s+(?P<name>\w+)", param)
+                        assert pm is not None
+                        params.append(StubMethod.Param(pm.group("name"), convert_type(pm.group("type")), const_mapping=const_mapping))
+                    else:
+                        pm = regex.fullmatch(r"(?P<type>\w+)\s+(?P<name>\w+)(?:\s*=\s*(?P<default>[^,]+))?", param)
+                        assert pm is not None
+                        params.append(StubMethod.Param(pm.group("name"), convert_type(pm.group("type")), pm.group("default"), const_mapping=const_mapping))
+
+                clazz.methods[name] = StubMethod(
+                    name,
+                    rtype,
+                    params,
+                    is_cdef,
+                    is_static="staticmethod" in decorators
+                )
+
+            decorators.clear()
+        # Decorator
+        elif m := regex.match(r"\s+@\s*(?P<decorator>\w[\w.]*)", line):
+            decorators.append(m.group("decorator"))
+        # DEF
+        elif m := regex.match(r"DEF\s+(?P<name>\w+)\s*=\s*(?P<value>.+)", line):
+            const_mapping[m.group("name")] = m.group("value")
+        else:
+            # if ("cdef" in line or "def" in line) and ":" in line:
+            #     print(f"Warning: ignored def or cdef line: {line}")
+            pass
+
+    if clazz is not None:
+        classes.append(clazz)
+
+    out_lines = [
+        "from typing import overload, Self, Any, Union",
+        ""
+    ]
+    for cls in classes:
+        print(f"gen_stub: generating class {cls.name}")
+        out_lines.extend(cls.stub())
+        out_lines.append("")
+    return "".join(f"{line}\n" for line in out_lines)
+
+
+if __name__ == '__main__':
+    result = gen_stub(open("output/_gdmath.pyx", encoding="utf8").read())
+    with open("output/_gdmath.pyi", "w") as f:
+        f.write(result)

codegen/templates/common_binary_and_inplace_op.pyx

@@ -15,11 +15,13 @@ cdef _VecClassName_ ___OpName___(self, _vTypeC_ other):
 
 #<OVERLOAD>
 cdef _VecClassName_ __i_OpName___(self, _VecClassName_ other):
+    #<RETURN_SELF>
     #<GEN>: gen_for_each_dim("self.{dim} _Op_= other.{dim}", _Dims_)
     return self
 
 #<OVERLOAD>
 cdef _VecClassName_ __i_OpName___(self, _vTypeC_ other):
+    #<RETURN_SELF>
     #<GEN>: gen_for_each_dim("self.{dim} _Op_= other", _Dims_)
     return self
 

codegen/templates/transform_2d.pyx

@@ -246,6 +246,7 @@ cdef class Transform2D:
         return t
 
     def __imatmul__(self, Transform2D other) -> Transform2D:
+        #<RETURN_SELF>
         cdef py_float xx = other.tdotx(self.xx, self.xy)
         cdef py_float xy = other.tdoty(self.xx, self.xy)
         cdef py_float yx = other.tdotx(self.yx, self.yy)
@@ -336,6 +337,7 @@ cdef class Transform2D:
 
 
     def translate_ip(self, Vec2 translation, /) -> Transform2D:
+        #<RETURN_SELF>
         self.ox = translation.x
         self.oy = translation.y
         return self
@@ -346,6 +348,7 @@ cdef class Transform2D:
         return t
 
     def rotate_ip(self, py_float rotation, /) -> Transform2D:
+        #<RETURN_SELF>
         self.__imatmul__(Transform2D.rotation(rotation))
         return self
 
@@ -355,6 +358,7 @@ cdef class Transform2D:
         return t
 
     def scale_ip(self, Vec2 scale, /) -> Transform2D:
+        #<RETURN_SELF>
         self.xx *= scale.x
         self.xy *= scale.y
         self.yx *= scale.x

codegen/templates/transform_3d.pyx

@@ -367,6 +367,7 @@ cdef class Transform3D:
         return t
 
     def __imatmul__(self, Transform3D other) -> Transform3D:
+        #<RETURN_SELF>
         cdef py_float xx = other.tdotx(self.xx, self.xy, self.xz)
         cdef py_float xy = other.tdoty(self.xx, self.xy, self.xz)
         cdef py_float xz = other.tdotz(self.xx, self.xy, self.xz)
@@ -417,6 +418,7 @@ cdef class Transform3D:
         return t
 
     def translate_ip(self, Vec3 translation, /) -> Transform3D:
+        #<RETURN_SELF>
         self.ox += translation.x
         self.oy += translation.y
         self.oz += translation.z
@@ -428,6 +430,7 @@ cdef class Transform3D:
         return t
 
     def rotate_ip(self, Vec3 axis, py_float angle, /) -> Transform3D:
+        #<RETURN_SELF>
         self.__imatmul__(Transform3D.rotating(axis, angle))
         return self
 
@@ -437,6 +440,7 @@ cdef class Transform3D:
         return t
 
     def scale_ip(self, Vec3 scale, /) -> Transform3D:
+        #<RETURN_SELF>
         self.xx *= scale.x
         self.yx *= scale.x
         self.zx *= scale.x