python/api2/demoSequencerDrawManager.py

import maya.api.OpenMayaUI as omui
import maya.api.OpenMaya as om

from maya import cmds

from numpy import *
import random

# Demo script to show how to use the sequencer custom
# drawing object API.
#
# This script can be loaded as a plugin or ran as a script

# Declare instances in global scope.
# Use globals().get() so this file can be safely reloaded without losing state
custom_sequencer_drawers = globals().get('custom_sequencer_drawers', None)

# Callback id used to cleanup on app exit (kMayaExiting). Stored globally so
# repeated runs/reloads don't accumulate callbacks.
maya_exit_callback_id = globals().get('maya_exit_callback_id', None)


class CustomSequencerDrawer:
    """
    Base class for drawing objects
    """

    def __init__(self):
        self.manager = omui.MSequencerCustomDrawManager()
        self.draw_ids = []
        self.init_draw_objects()

    def uninit(self):
        self.clear_primitives()

        for draw_id in self.draw_ids:
            self.manager.deregisterCustomDraw(draw_id)

    def create_timeline_draw_object(
        self, name, tooltip, draw_layer, bg_color=None
    ):
        """
        Creates and registers a custom timeline draw object with the sequencer
        draw manager.
        """
        draw_id = self.manager.registerCustomDrawTimeline(
            name,
            tooltip,
            draw_layer,  # Set a value lower than 1000 to draw behind ticks
        )

        if bg_color:
            self.manager.setBackgroundColor(draw_id, om.MColor(bg_color))

        self.draw_ids.append(draw_id)
        return draw_id

    def create_track_title_draw_object(
            self, name, tooltip, draw_layer, trackNumber,
            is_audio=False, bg_color=None
    ):
        """
        Creates and registers a custom track title draw object with the
        sequencer draw manager.
        """
        draw_id = self.manager.registerCustomDrawTrack(
            omui.MSequencerCustomDrawManager.kTrackTitle,
            name,
            tooltip,
            trackNumber,
            is_audio,
            draw_layer
        )

        self.set_track_properties(draw_id, bg_color)
        return draw_id

    def create_track_draw_object(
            self, name, tooltip, draw_layer, trackNumber,
            is_audio=False, bg_color=None
    ):
        """
        Creates and registers a custom track draw object with the sequencer
        draw manager.
        """
        draw_id = self.manager.registerCustomDrawTrack(
            omui.MSequencerCustomDrawManager.kTrack,
            name,
            tooltip,
            trackNumber,
            is_audio,
            draw_layer
        )

        self.set_track_properties(draw_id, bg_color)
        return draw_id

    def set_track_properties(
            self, draw_id, bg_color=None
    ):
        """
        Setup common attributes between track and track title draw objects
        """
        if bg_color:
            self.manager.setBackgroundColor(draw_id, om.MColor(bg_color))

        self.draw_ids.append(draw_id)
        return draw_id

    @classmethod
    def bracket_primitive(
        cls, start, stop, color, bottom, height, priority, label=''
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kBracket,
            om.MTime(start), om.MTime(stop), om.MColor(color),
            bottom, height, priority, label
        )

        return prim

    @classmethod
    def upper_outline_primitive(
        cls, start, stop, color, bottom, height, priority, label=''
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kUpperOutline,
            om.MTime(start), om.MTime(stop), om.MColor(color),
            bottom, height, priority, label
        )

        return prim

    @classmethod
    def full_outline_primitive(
        cls, start, stop, color, bottom, height, priority, label=''
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kFullOutline,
            om.MTime(start), om.MTime(stop), om.MColor(color),
            bottom, height, priority, label
        )

        return prim

    @classmethod
    def vertical_line_primitive(
        cls, time, color, bottom, height, priority, label=''
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kVerticalLine,
            om.MTime(time), om.MTime(time), om.MColor(color),
            bottom, height, priority, label
        )

        return prim

    @classmethod
    def filled_rect_primitive(
        cls, start, stop, color, bottom, height, priority, label=''
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kFilledRect,
            om.MTime(start), om.MTime(stop), om.MColor(color),
            bottom, height, priority, label
        )

        return prim
    
    @classmethod
    def text_primitive(
        cls, time, color, bottom, height, priority, label
    ):
        prim = omui.MSequencerDrawPrimitive(
            omui.MSequencerDrawPrimitive.kText,
            om.MTime(time), om.MTime(time), om.MColor(color),
            bottom, height, priority, label
        )

        return prim

    def init_draw_objects(self):
        pass

    def clear_primitives(self):
        manager = omui.MSequencerCustomDrawManager()

        for draw_id in self.draw_ids:
            manager.clearDrawPrimitives(draw_id)

        manager.requestSequencerRedraw()


class SequencerTimelineDrawDemo(CustomSequencerDrawer):
    """
    Demo custom drawing in the Sequencer timeline area.
    """
    def __init__(self):
        self.bg_draw_id = -1
        self.fg_draw_id = -1
        self.rgb_bg_color = [1.0, 1.0, 1.0]
        super().__init__()

    def init_draw_objects(self):
        manager = omui.MSequencerCustomDrawManager()
        self.init_bg_draw_object(manager)
        self.init_fg_draw_object(manager)

    def init_bg_draw_object(self, manager):
        """
        Setup draw object for sequencer timeline background
        This object will be drawn before/behind the ticks and frame numbers
        """
        draw_id = self.create_timeline_draw_object(
            'custom draw timeline 1',
            'My custom draw on timeline BG tooltip',
            draw_layer=5,  # Set a value lower than 1000 to draw behind ticks
            bg_color=[0.1, 0.3, 0.1, 0.25]
        )
        manager.setStartPrimitiveEditFunction(draw_id, self.startEdit)
        manager.setEditPrimitiveFunction(draw_id, self.editing)
        manager.setSetCopyPrimitivesFunction(draw_id, self.setCopying)
        manager.setStopPrimitiveEditFunction(draw_id, self.stopEdit)

        self.bg_draw_id = draw_id
        self.update_bg_primitives(manager)

    def init_fg_draw_object(self, manager):
        """
        Setup draw object for sequencer timeline foreground
        This object will be drawn over the ticks and frame numbers
        """
        self.fg_draw_id = self.create_timeline_draw_object(
            'custom draw timeline 2',
            '',
            draw_layer=1001,  # Set a value higher than 1000 to draw over ticks
        )

        self.update_fg_primitives(manager)

    def update_bg_primitives(self, manager):
        """
        Create primitives for background timeline draw object
        """
        prims = [
            CustomSequencerDrawer.full_outline_primitive(
                52, 615, [0.2, 0.8, 0.2],
                bottom=0.05, height=0.9, priority=1
            ),
            CustomSequencerDrawer.bracket_primitive(
                65, 175, [0.8, 0.2, 0.2],
                bottom=0.2, height=0.7, priority=10,
                label='Timeline BG Bracket'
            ),
            CustomSequencerDrawer.upper_outline_primitive(
                311, 353, [0.8, 0.8, 0.2],
                bottom=0.25, height=0.5, priority=1
            ),
            CustomSequencerDrawer.full_outline_primitive(
                405, 457, [0.8, 0.8, 0.8],
                bottom=0.1, height=0.75, priority=1
            ),
            CustomSequencerDrawer.vertical_line_primitive(
                515, [0.8, 0.1, 0.9],
                bottom=0, height=1, priority=1
            ),
        ]

        prim = CustomSequencerDrawer.filled_rect_primitive(
                210, 300, self.rgb_bg_color,
                bottom=0.1, height=0.75, priority=1
        )

        prim.tooltip = 'Filled Rect Tooltip'
        prim.setBuildMenuFunction(self.make_menu)
        prims.append(prim)

        manager.setDrawPrimitives(self.bg_draw_id, prims)
        manager.requestSequencerRedraw()

    def make_menu(self, parent_id, time):
        cmds.setParent(parent_id, menu=True)
        cmds.menuItem(label=f'Randomize color {time}',
                      command=self.randomize_bg_color)

    def randomize_bg_color(self, *args):
        print("Randomizing background color")
        self.rgb_bg_color = [random.random(), random.random(), random.random()]
        self.update_bg_primitives(omui.MSequencerCustomDrawManager())

    def update_fg_primitives(self, manager):
        """
        Create primitives for foreground timeline draw object
        """
        prims = [
            CustomSequencerDrawer.filled_rect_primitive(
                -100, -12, [0.8, 0.8, 0.2],
                bottom=0.25, height=0.5, priority=1
            ),
            CustomSequencerDrawer.full_outline_primitive(
                -210, -2, [0.8, 0.8, 0.8],
                bottom=0.1, height=0.75, priority=1
            ),
            CustomSequencerDrawer.text_primitive(
                -200, [1, 1, 1], bottom=0.1, height=0.75, priority=1,
                label='Timeline FG Draw Object'),
        ]

        manager.setDrawPrimitives(self.fg_draw_id, prims)

    def startEdit(self, time):
        print(f"Start editing at time {time}")

    def editing(self, time, currentTime):
        print(f"Editing at time {time}, currentTime={currentTime}")

    def stopEdit(self):
        print("Stop editing")

    def setCopying(self, copying):
        print(f"Set copying to {copying}")


class ShotTrackTitleDrawDemo(CustomSequencerDrawer):
    """
    Demo custom drawing in the Sequencer track title region for a shot track.
    """
    def __init__(self):
        self.toggled = True
        self.start = 10
        self.edit_delta = 0
        self.draw_id = 0
        super().__init__()

    def init_draw_objects(self):
        manager = omui.MSequencerCustomDrawManager()
        # Register custom draw for track title
        draw_id = self.create_track_title_draw_object(
            'custom draw track title',
            'My Custom Draw Track Title Tooltip',
            draw_layer=1,
            trackNumber=3,
            bg_color=[1.0, 0.3, 0.1, 0.5]
        )

        manager.setStartPrimitiveEditFunction(draw_id, self.startEdit)
        manager.setEditPrimitiveFunction(draw_id, self.editing)
        manager.setSetCopyPrimitivesFunction(draw_id, self.setCopying)
        manager.setStopPrimitiveEditFunction(draw_id, self.stopEdit)

        self.update_primitives(manager, draw_id)
        self.draw_id = draw_id

    def update_primitives(self, manager, draw_id):
        prims = [
            CustomSequencerDrawer.filled_rect_primitive(
                self.start + self.edit_delta,
                self.start + self.edit_delta + 80,
                [0.2, 0.8, 0.2, 0.75],
                bottom=0,
                height=0.5 if self.toggled else 0.75, 
                priority=1
            )
        ]
        prims[-1].tooltip = 'Track Title Filled Rect Tooltip Text'

        def sample_menu_item(self, *args):
            self.toggled = not self.toggled
            self.update_primitives(manager, draw_id)

        def make_menu(self, parent_id, time):
            cmds.setParent(parent_id, menu=True)
            cmds.menuItem(label='Toggle Height',
                          cb=self.toggled,
                          command=lambda *args: sample_menu_item(self, *args))

        prims[-1].setBuildMenuFunction(
            lambda parent_id, time: make_menu(self, parent_id, time))

        manager.setDrawPrimitives(draw_id, prims)
        manager.requestSequencerRedraw()

    def startEdit(self, x):
        print(f"Start editing track title at x {x}")
        self.edit_delta = 0

    def editing(self, delta, current_x):
        print(f"Editing track title at time {delta}, currentTime={current_x}")
        self.edit_delta = delta.value
        self.update_primitives(omui.MSequencerCustomDrawManager(), self.draw_id)

    def stopEdit(self):
        print("Stop editing track title")
        self.start += self.edit_delta
        self.edit_delta = 0
        self.update_primitives(omui.MSequencerCustomDrawManager(), self.draw_id)

    def setCopying(self, copying):
        print(f"Set copying to {copying} for track title")


class ShotTrackDrawDemo(CustomSequencerDrawer):
    """
    Demo custom drawing in the Sequencer for a shot track.
    """
    def __init__(self):
        self.edit_delta = None
        self.edit_copy = False
        self.fg_start = om.MTime(-60)
        self.fg_length = 100
        self.fg_draw_id = 0
        self.bg_draw_id = 0
        self.toggled = False
        self.fg_prims = []
        self.fg_edit_prim = None
        super().__init__()

    def init_draw_objects(self):
        manager = omui.MSequencerCustomDrawManager()
        # Register custom draw for track
        self.bg_draw_id = self.create_track_draw_object(
            'custom draw track BG',
            'My Custom Draw Track BG Tooltip',
            draw_layer=1,
            trackNumber=2,
            bg_color=[1.0, 0.3, 0.1, 0.5])

        self.fg_draw_id = self.create_track_draw_object(
            'custom draw track FG',
            'My Custom Draw Track FG Tooltip',
            draw_layer=1001,
            trackNumber=3,
            bg_color=[0.0, 0.3, 0.1, 0.5])

        manager.setStartPrimitiveEditFunction(self.fg_draw_id, self.startEdit)
        manager.setEditPrimitiveFunction(self.fg_draw_id, self.editing)
        manager.setSetCopyPrimitivesFunction(self.fg_draw_id, self.setCopying)
        manager.setStopPrimitiveEditFunction(self.fg_draw_id, self.stopEdit)

        self.fg_prims = [self.make_fg_obj_prims(self.fg_start)]

        self.update_bg_primitives(manager, self.bg_draw_id)
        self.update_fg_primitives(manager, self.fg_draw_id)

    def update_bg_primitives(self, manager, bg_draw_id):
        prims = [
            CustomSequencerDrawer.upper_outline_primitive(
                15, 245, [0.2, 0.2, 0.8],
                bottom=0.25 if self.toggled else 0.0,
                height=0.5 if self.toggled else 1.0,
                priority=1,
                label='Track BG Upper Outline'
            )
        ]
        prims[-1].tooltip = 'BG Upper Outline Tooltip Text'

        def make_menu(self, parent_id, time):
            cmds.setParent(parent_id, menu=True)
            cmds.menuItem(label='Toggle Mode',
                          checkBox=self.toggled,
                          command=self.do_toggle)

        prims[-1].setBuildMenuFunction(
            lambda parent_id, time: make_menu(self, parent_id, time))

        manager.setDrawPrimitives(bg_draw_id, prims)

    def make_fg_obj_prims(self, start, label=None):
        prim = CustomSequencerDrawer.full_outline_primitive(
            start.asUnits(om.MTime.kFilm),
            start.asUnits(om.MTime.kFilm) + self.fg_length,
            [0.8, 0.8 if self.toggled else 0.0, 0.2] if self.edit_delta is None else [1, 1, 1],
            bottom=0.1, height=0.7, priority=2,
            label='Drag to Move' if label is None else label
        )
        prim.tooltip = 'Move by dragging or copy by Ctrl + dragging, right click for more options'

        def make_menu(self, parent_id, time):
            cmds.setParent(parent_id, menu=True)
            cmds.menuItem(label='Toggle Mode',
                          checkBox=self.toggled,
                          command=self.do_toggle)
            cmds.menuItem(label='Duplicate',
                          command=lambda *args: self.duplicate(prim))
            cmds.menuItem(label='Delete',
                          command=lambda *args: self.delete_at(time),
                          enable=len(self.fg_prims) > 1 and self.find_fg_prim_at_time(time) is not None)

        prim.setBuildMenuFunction(
            lambda parent_id, time: make_menu(self, parent_id, time))

        return prim

    def update_fg_primitives(self, manager, fg_draw_id):
        prims = self.fg_prims.copy()

        if self.edit_delta and self.fg_edit_prim:
            self.fg_edit_prim.startTime = self.fg_start + self.edit_delta
            self.fg_edit_prim.endTime = self.fg_start + om.MTime(self.fg_length) + self.edit_delta

        if self.edit_copy and self.edit_delta:
            copy_prim = self.make_fg_obj_prims(self.fg_start, 'Ghost')
            copy_prim.color = om.MColor([0.8, 0.8, 0.2, 0.5])
            prims.append(copy_prim)

        manager.setDrawPrimitives(fg_draw_id, prims)

    def do_toggle(self, *args):
        self.toggled = not self.toggled

        for prim in self.fg_prims:
            prim.color = om.MColor([0.8, 0.8 if self.toggled else 0.0, 0.2])

        self.update_bg_primitives(omui.MSequencerCustomDrawManager(), self.bg_draw_id)
        self.update_fg_primitives(omui.MSequencerCustomDrawManager(), self.fg_draw_id)

    def find_fg_prim_at_time(self, time):
        for prim in self.fg_prims:
            if prim.startTime.asUnits(om.MTime.kFilm) <= time.asUnits(om.MTime.kFilm) <= prim.endTime.asUnits(om.MTime.kFilm):
                return prim
        return None

    def delete_at(self, time):
        print(f"Delete fg prim at time {time}")
        found = self.find_fg_prim_at_time(time)
        if found:
            self.fg_prims.remove(found)

        self.update_fg_primitives(omui.MSequencerCustomDrawManager(), self.fg_draw_id)

    def duplicate(self, prim):
        new_prim = self.make_fg_obj_prims(prim.startTime, 'Duplicate')
        new_prim.startTime = prim.endTime + om.MTime(10)
        new_prim.endTime = new_prim.startTime + om.MTime(self.fg_length)
        self.fg_prims.append(new_prim)
        self.update_fg_primitives(omui.MSequencerCustomDrawManager(), self.fg_draw_id)

    def startEdit(self, time):
        print(f"Start editing track at time {time}")
        self.fg_edit_prim = self.find_fg_prim_at_time(time)
        self.fg_start = self.fg_edit_prim.startTime if self.fg_edit_prim else time

    def editing(self, delta, currentTime):
        print(f"Editing track with delta {delta}, currentTime={currentTime}")
        self.edit_delta = delta
        self.update_fg_primitives(omui.MSequencerCustomDrawManager(), self.fg_draw_id)

    def stopEdit(self):
        print("Stop editing track")
        if self.edit_delta and self.fg_edit_prim:
            self.fg_edit_prim.startTime = self.fg_start + self.edit_delta
            self.fg_edit_prim.endTime = self.fg_start + om.MTime(self.fg_length) + self.edit_delta

        if self.edit_copy:
            self.fg_prims.insert(0, self.make_fg_obj_prims(self.fg_start, 'Copied'))

        self.edit_delta = None
        self.edit_copy = False
        self.fg_edit_prim = None
        self.update_fg_primitives(omui.MSequencerCustomDrawManager(), self.fg_draw_id)

    def setCopying(self, copying):
        print(f"Set copying to {copying} for track")
        self.edit_copy = copying


class AudioTrackDrawDemo(CustomSequencerDrawer):
    """
    Demo custom drawing in the Sequencer for an audio track.
    Includes both track and track title custom draw objects.
    """
    def __init__(self):
        self.track_draw_id = 0
        self.title_draw_id = 0
        self.track_number = 2
        self.toggled = False
        self.toggled_title = False
        super().__init__()

    def init_draw_objects(self):
        manager = omui.MSequencerCustomDrawManager()
        self.create_audio_track_draw_object(manager)
        self.create_audio_title_draw_object(manager)

    def create_audio_track_draw_object(self, manager):
        self.track_draw_id = self.create_track_draw_object(
            'custom draw audio track',
            'Audio Track Tooltip',
            draw_layer=2,
            trackNumber=self.track_number,
            is_audio=True,
            bg_color=[0.1, 0.8, 0.2, 0.7]
        )

        manager.setStartPrimitiveEditFunction(self.track_draw_id, self.startEdit)
        manager.setEditPrimitiveFunction(self.track_draw_id, self.editing)
        manager.setSetCopyPrimitivesFunction(self.track_draw_id, self.setCopying)
        manager.setStopPrimitiveEditFunction(self.track_draw_id, self.stopEdit)

        self.update_track_primitives(manager)

    def create_audio_title_draw_object(self, manager):
        self.title_draw_id = self.create_track_title_draw_object(
            'custom draw audio track title',
            'Audio Track Title Tooltip',
            draw_layer=1001,  # Set a value higher than default track title layer to draw over default title
            trackNumber=self.track_number,
            is_audio=True,
            bg_color=[0.8, 0.1, 0.2]
        )

        self.update_title_primitives(manager)

    def update_track_primitives(self, manager):
        prims = [
            CustomSequencerDrawer.filled_rect_primitive(
                100, 300 if self.toggled else 150, [0.2, 0.2, 0.2, 0.75],
                bottom=0.2, height=0.7, priority=1
            )
        ]
        prims[-1].tooltip = 'Audio Track Filled Rect Tooltip Text'

        def sample_menu_item(self, *args):
            self.toggled = not self.toggled
            self.update_track_primitives(manager)

        def make_menu(self, parent_id, time):
            cmds.setParent(parent_id, menu=True)
            cmds.menuItem(label=f'Example Audio Menu Item {time}',
                          command=lambda *args: sample_menu_item(self, *args))

        prims[-1].setBuildMenuFunction(
            lambda parent_id, time: make_menu(self, parent_id, time))

        manager.setDrawPrimitives(self.track_draw_id, prims)

    def update_title_primitives(self, manager):
        prims = [
            CustomSequencerDrawer.text_primitive(
                60, [1, 1, 1, 1],
                bottom=0.1, height=0.8, priority=1,
                label='Audio Track Title Object'
            ),
            CustomSequencerDrawer.filled_rect_primitive(
                10, 85, [0.2, 0.2, 0.2, 0.75],
                bottom=0.5 if self.toggled_title else 0.0,
                height=0.4 if self.toggled_title else 1.0,
                priority=1
            )
        ]
        prims[-1].tooltip = 'Audio Track Title Text Tooltip Text'

        def sample_menu_item(self, *args):
            self.toggled_title = not self.toggled_title
            self.update_title_primitives(manager)

        def make_menu(self, parent_id, time):
            cmds.setParent(parent_id, menu=True)
            cmds.menuItem(label=f'Example Audio Title Menu Item {time}',
                          command=lambda *args: sample_menu_item(self, *args))

        prims[-1].setBuildMenuFunction(
            lambda parent_id, time: make_menu(self, parent_id, time))

        manager.setDrawPrimitives(self.title_draw_id, prims)

    def startEdit(self, time):
        print(f"Start editing audio track obj at time {time}")

    def editing(self, time, currentTime):
        print(f"Editing audio track obj at time {time}, currentTime={currentTime}")

    def stopEdit(self):
        print("Stop editing audio track obj")

    def setCopying(self, copying):
        print(f"Set copying to {copying} for audio track obj")


def _setup_demo_sequencer_draw_manager():
    global custom_sequencer_drawers

    print("Demo sequencer draw manager")

    if custom_sequencer_drawers:
        for drawer in custom_sequencer_drawers:
            drawer.uninit()

    custom_sequencer_drawers = [
        SequencerTimelineDrawDemo(),
        ShotTrackTitleDrawDemo(),
        ShotTrackDrawDemo(),
        AudioTrackDrawDemo(),
    ]

    _setup_exit_callback()


def _setup_exit_callback():
    global maya_exit_callback_id

    if maya_exit_callback_id is not None:
        return

    maya_exit_callback_id = om.MSceneMessage.addCallback(
        om.MSceneMessage.kMayaExiting,
        _on_maya_exiting,
    )


def _cleanup_custom_sequencer_drawers():
    global custom_sequencer_drawers

    if not custom_sequencer_drawers:
        return

    for drawer in custom_sequencer_drawers:
        drawer.uninit()

    custom_sequencer_drawers = None


def _remove_exit_callback():
    global maya_exit_callback_id

    if maya_exit_callback_id is None:
        return

    try:
        om.MMessage.removeCallback(maya_exit_callback_id)
    except RuntimeError:
        # Callback may already be gone during shutdown.
        pass

    maya_exit_callback_id = None


def _on_maya_exiting(*args):
    _cleanup_custom_sequencer_drawers()
    _remove_exit_callback()


def initializePlugin(mobject):
    _setup_exit_callback()
    _setup_demo_sequencer_draw_manager()


def uninitializePlugin(mobject):
    global custom_sequencer_drawers

    print("Uninitialize plugin and clear custom sequencer drawers")

    _cleanup_custom_sequencer_drawers()
    _remove_exit_callback()


def maya_useNewAPI():
    pass


if __name__ == '__main__':
    _setup_demo_sequencer_draw_manager()