Menu context manager for Blender.

blender_menus.py

"""Menu context manager for Blender.

Example:
    >>> with Menu('Custom Menu') as menu:
    ...     with menu.add_submenu('Submenu') as submenu:
    ...         submenu.add_operator('mesh.primitive_cube_add')
    ...     menu.add_operator(lambda: 1/0, 'Raise Exception')
    >>> menu.register()
    >>> menu.unregister()
"""

from __future__ import annotations

__all__ = ['Menu']

from abc import ABC, abstractmethod
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING, Callable, List, Optional
from typing import Set, Type, Union, Generator, cast
from uuid import uuid4
import weakref

import bpy

import inflection


class AbstractMenu(ABC):
    """Base class for dynamically generating Blender menus.

    Each menu gets a unique ID, and is only populated when registered.
    This means that the menu can be unloaded and modified.
    """

    __slots__ = [
        '_id', '_name', '_parent', '_operator_ids',
        '_custom_menu', '_custom_operators', '_custom_submenus',
    ]

    def __init__(self, name: str, parent: Optional[AbstractMenu] = None) -> None:
        self._id = f'TOPBAR_MT_custom_menu_{uuid4().hex}'
        self._name = name

        self._parent: Optional[weakref.ReferenceType[AbstractMenu]]
        if parent is None:
            self._parent = None
        else:
            self._parent = weakref.proxy(parent)

        self._operator_ids: List[str] = []
        self._custom_menu: Optional[Type[bpy.types.Menu]] = None
        self._custom_operators: List[Type[bpy.types.Operator]] = []
        self._custom_submenus: List[AbstractMenu] = []

    def __enter__(self) -> AbstractMenu:
        return self

    def __exit__(self, *args) -> bool:
        return not any(args)

    @property
    def registered(self) -> bool:
        """Determine if the class has been registered."""
        return self._custom_menu is not None

    @property
    def id(self) -> str:
        """Get the unique menu ID."""
        return self._id

    @property
    def name(self) -> str:
        """Get the menu name."""
        return self._name

    @name.setter
    def name(self, name: str) -> None:
        """Set a new menu name."""
        if self.registered:
            raise RuntimeError('unable to rename while the menu is registered')
        self._name = name

    @property
    def parent(self) -> Optional[weakref.ReferenceType[AbstractMenu]]:
        """Get the menu parent."""
        return self._parent

    @property
    def operators(self) -> Generator[str, None, None]:
        """Get the operators."""
        yield from self._operator_ids

    @property
    def submenus(self) -> Generator[AbstractMenu, None, None]:
        """Get the submenus."""
        yield from self._custom_submenus

    @abstractmethod
    def _menu_init(self) -> Type[bpy.types.Menu]:
        """Create a new custom menu class."""

    def _build_layout(self, layout: bpy.types.UILayout) -> None:
        """Build the layout for a menu.
        This is for the `bpy.types.Menu.draw` method.
        """
        for menu in self._custom_submenus:
            layout.menu(menu.id)
        for operator in self._operator_ids:
            layout.operator(operator)

    def register(self) -> None:
        """Register all custom classes to Blender."""
        if self.registered:
            raise RuntimeError('menu is already registered')

        # Create a new menu with all the required parameters
        self._custom_menu = self._menu_init()
        self._custom_menu.__name__ = self.id

        # Register all the classes
        bpy.utils.register_class(self._custom_menu)

        for operator in self._custom_operators:
            bpy.utils.register_class(operator)

        for submenu in self._custom_submenus:
            submenu.register()

    def unregister(self) -> None:
        """Unregister all custom classes from Blender."""
        if not self.registered:
            raise RuntimeError('menu is not yet registered')

        # Unregister all the classes
        for submenu in reversed(self._custom_submenus):
            submenu.unregister()

        for operator in reversed(self._custom_operators):
            with suppress(RuntimeError):
                bpy.utils.unregister_class(operator)

        with suppress(RuntimeError):
            bpy.utils.unregister_class(self._custom_menu)

        self._custom_menu = None

    def add_operator(self, operator: Union[str, Callable], name: Optional[str] = None) -> str:
        """Add an operator to the menu.
        If a function is given, then a new operator will be created to
        wrap that function.

        Example:
            >>> with Menu() as menu:
            ...     menu.add_operator('mesh.primitive_cube_add')
            ...     menu.add_operator(lambda: print(5), name='Print 5')
        """
        if self.registered:
            raise RuntimeError('unable to add operators while the menu is registered')

        # Create an operator to wrap the function
        if callable(operator):
            fn = operator
            if name is None:
                name = inflection.titleize(fn.__name__)

            class CustomOperator(bpy.types.Operator):
                """Create an operator for a function."""

                bl_idname = 'wm.' + uuid4().hex
                bl_label = name

                def execute(self, context: bpy.types.Context) -> Set[str]:
                    fn()
                    return {'FINISHED'}

            # Store the class for registering/unregistering
            self._custom_operators.append(CustomOperator)
            operator = CustomOperator.bl_idname

        # Store the ID for when building the layouts
        self._operator_ids.append(operator)
        return operator

    @contextmanager
    def add_submenu(self, name: str) -> Generator[AbstractMenu, None, None]:
        """Create a new submenu.

        Example:
            >>> with Menu() as menu:
            ...     with menu.add_submenu() as submenu:
            ...         pass
        """
        if self.registered:
            raise RuntimeError('unable to add submenus while the menu is registered')

        with Submenu(name, parent=self) as submenu:
            yield submenu
        self._custom_submenus.append(submenu)


class Menu(AbstractMenu):
    """Dynamically generate a new top menu."""

    def _menu_init(self) -> Type[bpy.types.Menu]:
        """Create a new custom menu class."""
        parent = self

        class CustomTopMenu(bpy.types.Menu):
            """Main class for the menu."""
            bl_label = self.name

            def draw(self, context: bpy.types.Context) -> None:
                """Draw the layout."""
                parent._build_layout(self.layout)

            def menu_draw(self, context: bpy.types.Context) -> None:
                """Add to the top bar."""
                self.layout.menu(parent.id)

        return CustomTopMenu

    def register(self) -> None:
        """Register the classes and add the menu to the top bar."""
        super().register()

        if TYPE_CHECKING:
            self._custom_menu = cast(Type[bpy.types.Menu], self._custom_menu)
        bpy.types.TOPBAR_MT_editor_menus.append(self._custom_menu.menu_draw)

    def unregister(self) -> None:
        """Unregister the classes and remove the menu from the top bar."""
        menu = self._custom_menu
        super().unregister()

        if TYPE_CHECKING:
            menu = cast(Type[bpy.types.Menu], menu)
        bpy.types.TOPBAR_MT_editor_menus.remove(menu.menu_draw)


class Submenu(AbstractMenu):
    """Dynamically generate submenus.

    Note that a submenu may be unregistered separately to its parent
    menu, but it must be registered again or errors will occur.
    """

    def _menu_init(self) -> Type[bpy.types.Menu]:
        """Create a new custom menu class."""
        parent = self

        class CustomSubmenu(bpy.types.Menu):
            """Main class for the menu."""

            bl_label = self.name

            def draw(self, context: bpy.types.Context) -> None:
                """Draw the layout."""
                parent._build_layout(self.layout)

        return CustomSubmenu


def example() -> AbstractMenu:
    """Example/test generating a basic menu."""

    def print_time():
        """Print the current time.
        The function name will get automatically converted to "Print Time".
        """
        import time
        print(time.time())

    with Menu('Custom Menu') as menu:
        with menu.add_submenu('Submenu') as submenu:
            submenu.add_operator('mesh.primitive_cube_add')
            submenu.add_operator(print_time)
        menu.add_operator(lambda: 1/0, 'Raise Exception')

    menu.register()
    return menu


if __name__ == '__main__':
    example()