drawing arrows

Author: qwiglydee

src/ifcblenderexport/blenderbim/bim/decoration.py

@@ -1,10 +1,15 @@
 """Viewport decorations"""
 import math
+from functools import reduce
 from bpy.types import SpaceView3D
 from mathutils import Vector
 import bpy
 import blf
 from bpy_extras.view3d_utils import location_3d_to_region_2d
+import gpu
+import bgl
+from gpu.types import GPUShader, GPUBatch, GPUIndexBuf, GPUVertBuf, GPUVertFormat
+from gpu_extras.batch import batch_for_shader
 
 
 class ViewDecorator(object):
@@ -31,11 +36,81 @@ class DimensionDecorator(ViewDecorator):
     - puts metric text next to each segment
     """
 
+    VERT_GLSL = """
+    uniform mat4 viewMatrix;
+    in vec3 pos;
+    out vec4 gl_Position;
+
+    void main() {
+        gl_Position = viewMatrix * vec4(pos, 1.0);
+    }
+    """
+    GEOM_GLSL = """
+    layout(lines) in;
+    layout(line_strip, max_vertices=8) out;
+
+    uniform float angle;
+    uniform float length;
+
+    void main() {
+        vec4 p0 = gl_in[0].gl_Position, p1 = gl_in[1].gl_Position;
+        float c = cos(angle), s = sin(angle);
+        mat4 rot_a = mat4( c, -s, 0, 0,
+                          +s,  c, 0, 0,
+                           0,  0, 1, 0,
+                           0,  0, 0, 1);
+        mat4 rot_b = mat4( c, +s, 0, 0,
+                          -s,  c, 0, 0,
+                           0,  0, 1, 0,
+                           0,  0, 0, 1);
+
+        vec4 dir = normalize(p1 - p0);
+
+        vec4 arr_a = rot_a * dir * length;
+        vec4 arr_b = rot_b * dir * length;
+
+        gl_Position = p0;
+        EmitVertex();
+        gl_Position = p1;
+        EmitVertex();
+        EndPrimitive();
+
+        gl_Position = p0 + arr_a;
+        EmitVertex();
+        gl_Position = p0;
+        EmitVertex();
+        gl_Position = p0 + arr_b;
+        EmitVertex();
+        EndPrimitive();
+
+        gl_Position = p1 - arr_a;
+        EmitVertex();
+        gl_Position = p1;
+        EmitVertex();
+        gl_Position = p1 - arr_b;
+        EmitVertex();
+        EndPrimitive();
+    }
+    """
+    FRAG_GLSL = """
+    uniform vec3 color;
+    out vec4 fragColor;
+
+    void main() {
+        fragColor = vec4(color, 1.0);
+    }
+    """
+
     def __init__(self, props, context):
         self.context = context
         self.props = props
         self.font_id = 0  # TODO: take font from styles
         self.dpi = context.preferences.system.dpi
+        self.shader = self.create_shader()
+
+    @classmethod
+    def create_shader(cls):
+        return GPUShader(vertexcode=cls.VERT_GLSL, fragcode=cls.FRAG_GLSL, geocode=cls.GEOM_GLSL)
 
     def __call__(self):
         # get active drawing, if any
@@ -48,9 +123,10 @@ def __call__(self):
             return
         curve = collection.all_objects['IfcAnnotation/Dimension']
 
-        for segm in self.iter_segments(curve):
+        segments = list(self.iter_segments(curve))
+        self.draw_arrows(segments)
+        for segm in segments:
             self.draw_label(segm)
-            self.draw_arrow(segm)
 
     def iter_segments(self, curve):
         """Yields each segment converted to world coords
@@ -105,5 +181,18 @@ def draw_label(self, segm):
         blf.draw(self.font_id, text)
         blf.disable(self.font_id, blf.ROTATION)
 
-    def draw_arrow(self, segm):
-        pass
+    def draw_arrows(self, segments):
+        def coords(segm):
+            return [(segm[0].x, segm[0].y, segm[0].z),
+                    (segm[1].x, segm[1].y, segm[1].z)]
+        points = list(reduce(lambda points, segm: points + coords(segm),
+                             segments, []))
+        batch = batch_for_shader(self.shader, 'LINES', {'pos': points})
+        self.shader.bind()
+        matrix = self.context.region_data.perspective_matrix
+        self.shader.uniform_float("viewMatrix", matrix)
+        # TODO: get everything from styles
+        self.shader.uniform_float('color', (1.0, 1.0, 1.0))
+        self.shader.uniform_float('angle', math.pi / 12)
+        self.shader.uniform_float('length', 0.05)
+        batch.draw(self.shader)