storypencil/dopesheet_overlay.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import typing
   4
   5 import bpy
   6 import gpu
   7 from gpu_extras.batch import batch_for_shader
   8
   9 from .utils import (redraw_all_areas_by_type)
  10 from .synchro import (is_secondary_window, window_id, get_main_strip)
  11
  12 Int3 = typing.Tuple[int, int, int]
  13
  14 Float2 = typing.Tuple[float, float]
  15 Float3 = typing.Tuple[float, float, float]
  16 Float4 = typing.Tuple[float, float, float, float]
  17
  18
  19 class LineDrawer:
  20     def __init__(self):
  21         self._format = gpu.types.GPUVertFormat()
  22         self._pos_id = self._format.attr_add(
  23             id="pos", comp_type="F32", len=2, fetch_mode="FLOAT"
  24         )
  25         self._color_id = self._format.attr_add(
  26             id="color", comp_type="F32", len=4, fetch_mode="FLOAT"
  27         )
  28
  29         self.shader = gpu.shader.from_builtin('2D_UNIFORM_COLOR')
  30
  31     def draw(
  32             self,
  33             coords: typing.List[Float2],
  34             indices: typing.List[Int3],
  35             color: Float4,
  36     ):
  37         if not coords:
  38             return
  39
  40         gpu.state.blend_set('ALPHA')
  41
  42         self.shader.uniform_float("color", color)
  43
  44         batch = batch_for_shader(self.shader, 'TRIS', {"pos": coords}, indices=indices)
  45         batch.program_set(self.shader)
  46         batch.draw()
  47
  48         gpu.state.blend_set('NONE')
  49
  50
  51 def get_scene_strip_in_out(strip):
  52     """ Return the in and out keyframe of the given strip in the scene time reference"""
  53     shot_in = strip.scene.frame_start + strip.frame_offset_start
  54     shot_out = shot_in + strip.frame_final_duration - 1
  55     return (shot_in, shot_out)
  56
  57
  58 def draw_callback_px(line_drawer: LineDrawer):
  59     context = bpy.context
  60     region = context.region
  61     main_scene = context.scene.storypencil_main_scene
  62     if main_scene is None:
  63         return
  64
  65     use_win = main_scene.storypencil_use_new_window
  66     wm = context.window_manager
  67
  68     if (
  69             (use_win and not wm.storypencil_settings.active)
  70             or not wm.storypencil_settings.show_main_strip_range
  71             or (use_win and not is_secondary_window(wm, window_id(context.window)))
  72             or (not use_win and context.scene == main_scene)
  73     ):
  74         return
  75
  76     # get main strip driving the sync
  77     strip = get_main_strip(wm)
  78
  79     if not strip or strip.scene != context.scene:
  80         return
  81
  82     xwin1, ywin1 = region.view2d.region_to_view(0, 0)
  83     one_pixel_further_x = region.view2d.region_to_view(1, 1)[0]
  84     pixel_size_x = one_pixel_further_x - xwin1
  85     rect_width = 1
  86
  87     shot_in, shot_out = get_scene_strip_in_out(strip)
  88     key_coords_in = [
  89         (
  90             shot_in - rect_width * pixel_size_x,
  91             ywin1,
  92         ),
  93         (
  94             shot_in + rect_width * pixel_size_x,
  95             ywin1,
  96         ),
  97         (
  98             shot_in + rect_width * pixel_size_x,
  99             ywin1 + context.region.height,
 100         ),
 101         (
 102             shot_in - rect_width * pixel_size_x,
 103             ywin1 + context.region.height,
 104         ),
 105     ]
 106
 107     key_coords_out = [
 108         (
 109             shot_out - rect_width * pixel_size_x,
 110             ywin1,
 111         ),
 112         (
 113             shot_out + rect_width * pixel_size_x,
 114             ywin1,
 115         ),
 116         (
 117             shot_out + rect_width * pixel_size_x,
 118             ywin1 + context.region.height,
 119         ),
 120         (
 121             shot_out - rect_width * pixel_size_x,
 122             ywin1 + context.region.height,
 123         ),
 124     ]
 125
 126     indices = [(0, 1, 2), (2, 0, 3)]
 127     # Draw the IN frame in green
 128     # hack: in certain cases, opengl draw state is invalid for the first drawn item
 129     #       resulting in a non-colored line
 130     #       => draw it a first time with a null alpha, so that the second one is drawn correctly
 131     line_drawer.draw(key_coords_in, indices, (0, 0, 0, 0))
 132     line_drawer.draw(key_coords_in, indices, (0.3, 0.99, 0.4, 0.5))
 133     # Draw the OUT frame un red
 134     line_drawer.draw(key_coords_out, indices, (0.99, 0.3, 0.4, 0.5))
 135
 136
 137 def tag_redraw_all_dopesheets():
 138     redraw_all_areas_by_type(bpy.context, 'DOPESHEET')
 139
 140
 141 # This is a list so it can be changed instead of set
 142 # if it is only changed, it does not have to be declared as a global everywhere
 143 cb_handle = []
 144
 145
 146 def callback_enable():
 147     if cb_handle:
 148         return
 149
 150     # Doing GPU stuff in the background crashes Blender, so let's not.
 151     if bpy.app.background:
 152         return
 153
 154     line_drawer = LineDrawer()
 155     # POST_VIEW allow to work in time coordinate (1 unit = 1 frame)
 156     cb_handle[:] = (
 157         bpy.types.SpaceDopeSheetEditor.draw_handler_add(
 158             draw_callback_px, (line_drawer,), 'WINDOW', 'POST_VIEW'
 159         ),
 160     )
 161
 162     tag_redraw_all_dopesheets()
 163
 164
 165 def callback_disable():
 166     if not cb_handle:
 167         return
 168
 169     try:
 170         bpy.types.SpaceDopeSheetEditor.draw_handler_remove(cb_handle[0], 'WINDOW')
 171     except ValueError:
 172         # Thrown when already removed.
 173         pass
 174     cb_handle.clear()
 175
 176     tag_redraw_all_dopesheets()
 177
 178
 179 def register():
 180     callback_enable()
 181
 182
 183 def unregister():
 184     callback_disable()