Move input helper into WebGPU class

webgpu/gpu.py

@@ -1,12 +1,29 @@
 import js
+import sys
 
 from .colormap import Colormap
 from .uniforms import Uniforms
 from .utils import to_js
+from .input_handler import InputHandler
 
+async def init_webgpu(canvas):
+    """Initialize WebGPU, create device and canvas"""
+    if not js.navigator.gpu:
+        js.alert("WebGPU is not supported")
+        sys.exit(1)
+
+    adapter = await js.navigator.gpu.requestAdapter()
+
+    if not adapter:
+        js.alert("WebGPU is not supported")
+        sys.exit(1)
+
+    device = await adapter.requestDevice()
+
+    return WebGPU(device, canvas)
 
 class WebGPU:
-    """WebGPU management class, handles "global" state, like device, canvas, colormap and uniforms"""
+    """WebGPU management class, handles "global" state, like device, canvas, frame/depth buffer, colormap and uniforms"""
 
     def __init__(self, device, canvas):
         self.render_function = None
@@ -43,6 +60,8 @@ def __init__(self, device, canvas):
                 }
             )
         )
+        self.input_handler = InputHandler(canvas, self.uniforms)
+
 
     def begin_render_pass(self, command_encoder):
         render_pass_encoder = command_encoder.beginRenderPass(
@@ -76,3 +95,7 @@ def __del__(self):
         self.depth_texture.destroy()
         del self.uniforms
         del self.colormap
+
+        # unregister is needed to remove circular references
+        self.input_handler.unregister_callbacks()
+        del self.input_handler

webgpu/input_handler.py

@@ -4,8 +4,10 @@
 
 
 class InputHandler:
-    def __init__(self, gpu):
-        self.gpu = gpu
+    def __init__(self, canvas, uniforms, render_function=None):
+        self.canvas = canvas
+        self.uniforms = uniforms
+        self.render_function = render_function
         self._is_moving = False
 
         self._callbacks = {}
@@ -20,14 +22,15 @@ def on_mouseup(self, _):
 
     def on_mousemove(self, ev):
         if self._is_moving:
-            self.gpu.uniforms.mat[12] += ev.movementX / self.gpu.canvas.width * 1.8
-            self.gpu.uniforms.mat[13] -= ev.movementY / self.gpu.canvas.height * 1.8
-            js.requestAnimationFrame(self.gpu.render_function)
+            self.uniforms.mat[12] += ev.movementX / self.canvas.width * 1.8
+            self.uniforms.mat[13] -= ev.movementY / self.canvas.height * 1.8
+            if self.render_function:
+                js.requestAnimationFrame(self.render_function)
 
     def unregister_callbacks(self):
         for event in self._callbacks:
             for func in self._callbacks[event]:
-                self.gpu.canvas.removeEventListener(event, func)
+                self.canvas.removeEventListener(event, func)
                 func.destroy()
         self._callbacks = {}
 
@@ -36,7 +39,7 @@ def on(self, event, func):
         if event not in self._callbacks:
             self._callbacks[event] = []
         self._callbacks[event].append(func)
-        self.gpu.canvas.addEventListener(event, func)
+        self.canvas.addEventListener(event, func)
 
     def register_callbacks(self):
         self.unregister_callbacks()
@@ -46,3 +49,7 @@ def register_callbacks(self):
 
     def __del__(self):
         self.unregister_callbacks()
+        if self.render_function:
+            js.cancelAnimationFrame(self.render_function)
+            self.render_function.destroy()
+            self.render_function = None

webgpu/main.py

@@ -1,35 +1,19 @@
-import sys
+"""Main file for the webgpu example, creates a small 2d mesh and renders it using WebGPU"""
 
 import js
 from pyodide.ffi import create_proxy
 
-from .gpu import WebGPU
-from .input_handler import InputHandler
+from .gpu import init_webgpu
 from .mesh import MeshRenderObject
 
 gpu = None
-input_handler = None
 mesh_object = None
 
 
-async def main(canvas=None):
-    global gpu, input_handler, mesh_object
+async def main():
+    global gpu, mesh_object
 
-    if not js.navigator.gpu:
-        js.alert("WebGPU is not supported")
-        sys.exit(1)
-
-    canvas = canvas or js.document.getElementById("canvas")
-    adapter = await js.navigator.gpu.requestAdapter()
-
-    if not adapter:
-        js.alert("WebGPU is not supported")
-        sys.exit(1)
-
-    device = await adapter.requestDevice()
-
-    gpu = WebGPU(device, canvas)
-    input_handler = InputHandler(gpu)
+    gpu = await init_webgpu(js.document.getElementById("canvas"))
 
     from netgen.occ import unit_square
 
@@ -52,28 +36,27 @@ def render(time):
         # copy camera position etc. to GPU
         gpu.uniforms.update_buffer()
 
-        command_encoder = device.createCommandEncoder()
+        command_encoder = gpu.device.createCommandEncoder()
 
         render_pass_encoder = gpu.begin_render_pass(command_encoder)
         mesh_object.draw(render_pass_encoder)
         render_pass_encoder.end()
 
-        device.queue.submit([command_encoder.finish()])
+        gpu.device.queue.submit([command_encoder.finish()])
 
-    gpu.render_function = create_proxy(render)
+    render_function = create_proxy(render)
+    gpu.input_handler.render_function = render_function
 
-    js.requestAnimationFrame(gpu.render_function)
+    js.requestAnimationFrame(render_function)
 
 
 def cleanup():
     print("cleanup")
-    global gpu, input_handler, mesh_object
-    if "input_handler" in globals():
-        if input_handler is not None:
-            input_handler.unregister_callbacks()
-        del mesh_object
-        del input_handler
+    global gpu, mesh_object
+    if "gpu" in globals():
         del gpu
+    if "mesh_object" in globals():
+        del mesh_object
 
 
 def reload_package(package_name):

webgpu/uniforms.py

@@ -29,6 +29,8 @@ class ColormapUniform(ct.Structure):
 
 
 class Uniforms(ct.Structure):
+    """Uniforms class, derived from ctypes.Structure to ensure correct memory layout"""
+
     _fields_ = [
         ("mat", ct.c_float * 16),
         ("clipping_plane", ClippingPlaneUniform),
@@ -85,6 +87,7 @@ def get_binding(self):
         return [{"binding": Binding.UNIFORMS, "resource": {"buffer": self.buffer}}]
 
     def update_buffer(self):
+        """Copy the current data to the GPU buffer"""
         data = js.Uint8Array.new(bytes(self))
         self.device.queue.writeBuffer(self.buffer, 0, data)