storypencil/dopesheet_overlay.py

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