imgui bindings: add stubs for ImVec2/4 math operators (fix #267)

pthom · Oct 20, 2024 · 97ee7cc · 97ee7cc
1 parent c70ddce
commit 97ee7cc
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 10 deletions.
diff --git a/bindings/imgui_bundle/imgui/__init__.pyi b/bindings/imgui_bundle/imgui/__init__.pyi
@@ -10,6 +10,9 @@ from typing import (
     overload,
     Iterator,
     Callable,
+    Union,
+    Protocol,
+    TypeVar,
 )
 import numpy as np
 import enum
@@ -321,6 +324,16 @@ VERTEX_BUFFER_UV_OFFSET: int
 VERTEX_BUFFER_COL_OFFSET: int
 INDEX_SIZE: int
 
+# VecProtocol: add __add__, __sub__, __mul__, __truediv__, __neg__ to ImVec2 and ImVec4
+TVec = TypeVar("TVec", bound="VecProtocol")
+
+class VecProtocol(Protocol[TVec]):
+    def __add__(self: TVec, other: TVec) -> TVec: ...
+    def __sub__(self: TVec, other: TVec) -> TVec: ...
+    def __mul__(self: TVec, other: Union[TVec, float]) -> TVec: ...
+    def __truediv__(self: TVec, other: Union[TVec, float]) -> TVec: ...
+    def __neg__(self: TVec) -> TVec: ...
+
 ##################################################
 #    AUTO GENERATED CODE BELOW
 ##################################################
@@ -469,7 +482,7 @@ INDEX_SIZE: int
 # #endif
 #
 
-class ImVec2:
+class ImVec2(VecProtocol["ImVec2"]):
     # float                                   x,     /* original C++ signature */
     x: float
     # y;    /* original C++ signature */
@@ -509,7 +522,7 @@ class ImVec2:
     # #endif
     #
 
-class ImVec4:
+class ImVec4(VecProtocol["ImVec4"]):
     """ImVec4: 4D vector used to store clipping rectangles, colors etc. [Compile-time configurable type]"""
 
     # float                                                     x,     /* original C++ signature */

diff --git a/external/imgui/bindings/litgen_options_imgui.py b/external/imgui/bindings/litgen_options_imgui.py
@@ -179,14 +179,6 @@ def preprocess_ImGuiTestGuiFunc(code: str) -> str:
     options.member_exclude_by_type__regex += "|^ImMovingAverage|^Str$|^ImGuiPerfTool|^ImGuiCaptureToolUI|^ImGuiCaptureContext|^ImGuiCaptureArgs"
     options.fn_exclude_by_param_type__regex += "|^ImGuiCaptureArgs"
 
-    def postprocess_stub(code: str):
-        # any function that accepts a TestRef param should also accept str (which is convertible to TestRef)
-        r = code.replace(": TestRef", ": Union[TestRef, str]")
-        r = r.replace("(TestEngineExportFormat)0", "TestEngineExportFormat.j_unit_xml")
-        return r
-
-    options.postprocess_stub_function = postprocess_stub
-
 
 def litgen_options_imgui(
     options_type: ImguiOptionsType, docking_branch: bool
@@ -441,6 +433,26 @@ def litgen_options_imgui(
     elif options_type == ImguiOptionsType.imgui_test_engine:
         add_imgui_test_engine_options(options)
 
+    def postprocess_stub_add_vec_protocol(stub_code: str) -> str:
+        stub_code = stub_code.replace(
+            "class ImVec2:", "class ImVec2(VecProtocol['ImVec2']):"
+        )
+        stub_code = stub_code.replace(
+            "class ImVec4:", "class ImVec4(VecProtocol['ImVec4']):"
+        )
+        return stub_code
+
+    def postprocess_stub_test_engine(code: str) -> str:
+        # any function that accepts a TestRef param should also accept str (which is convertible to TestRef)
+        r = code.replace(": TestRef", ": Union[TestRef, str]")
+        r = r.replace("(TestEngineExportFormat)0", "TestEngineExportFormat.j_unit_xml")
+        return r
+
+    if options_type == ImguiOptionsType.imgui_test_engine:
+        options.postprocess_stub_function = postprocess_stub_test_engine
+    else:
+        options.postprocess_stub_function = postprocess_stub_add_vec_protocol
+
     return options