c4d_treeview.pyp

# Example of using the Cinema 4D Tree View GUI in Python.
# Originally written by Niklas Rosenstein
# Modifications by Donovan Keith

import c4d
import os
import weakref

# Be sure to use a unique ID obtained from http://www.plugincafe.com/.
PLUGIN_ID = 9912399
PLUGIN_NAME = "Python TreeView Example"
PLUGIN_HELP = "Show the current scene hierarchy in a tree."

# Bit for the Checkbox in the GUI. Kind of a bad practice to use this
# but reduces complexity for this example.
BIT_CHECKED = 256

# TreeView Column IDs.
ID_CHECKBOX = 1
ID_ICON = 2
ID_TREEVIEW = 3


class Hierarchy(c4d.gui.TreeViewFunctions):
    def __init__(self, dlg):
        # Avoid a cyclic reference.
        self._dlg = weakref.ref(dlg)

    def GetFirst(self, root, userdata):
        """
        Return the first element in the hierarchy. This can be any Python
        object. With Cinema 4D nodes, it is easy to implement, but it could
        as well be a Python integer which you can use as an index in a list.

        Returns:
          (any or None): The first element in the hierarchy or None.
        """

        doc = c4d.documents.GetActiveDocument()
        return doc.GetFirstObject()

    def GetDown(self, root, userdata, obj):
        """
        Returns:
          (any or None): The child element going from *obj* or None.
        """

        return obj.GetDown()

    def GetNext(self, root, userdata, obj):
        """
        Returns:
          (any or None): The next element going from *obj* or None.
        """

        return obj.GetNext()

    def GetPred(self, root, userdata, obj):
        """
        Returns:
          (any or None): The previous element going from *obj* or None.
        """

        return obj.GetPred()

    def GetName(self, root, userdata, obj):
        """
        Returns:
          (str): The name to display for *obj*.
        """

        return obj.GetName()

    def Open(self, root, userdata, obj, opened):
        """
        Called when the user opens or closes the children of *obj* in
        the Tree View.
        """

        doc = c4d.documents.GetActiveDocument()
        doc.StartUndo()
        doc.AddUndo(c4d.UNDOTYPE_CHANGE_SELECTION, obj)

        if opened:
            obj.SetBit(c4d.BIT_OFOLD)
        else:
            obj.DelBit(c4d.BIT_OFOLD)
        doc.EndUndo()
        c4d.EventAdd()

    def IsOpened(self, root, userdata, obj):
        """
        Returns:
          (bool): True if *obj* is opened (expaneded), False if it is
            closed (collapsed).
        """

        return obj.GetBit(c4d.BIT_OFOLD)

    def DeletePressed(self, root, userdata):
        """
        Called when the user right-click Deletes or presses the Delete key.
        Should delete all selected elements.
        """

        def delete_recursive(op):
            if op is None: return
            if op.GetBit(c4d.BIT_ACTIVE) == True:
                doc.AddUndo(c4d.UNDOTYPE_DELETE, op)
                op.Remove()
                return
            for child in op.GetChildren():
                delete_recursive(child)

        doc = c4d.documents.GetActiveDocument()
        doc.StartUndo()
        for obj in doc.GetObjects():
            delete_recursive(obj)
        doc.EndUndo()
        c4d.EventAdd()

    def Select(self, root, userdata, obj, mode):
        """
        Called when the user selects an element.
        """

        doc = c4d.documents.GetActiveDocument()
        doc.StartUndo()

        doc.AddUndo(c4d.UNDOTYPE_CHANGE_SMALL, obj)
        doc.SetActiveObject(obj, mode)

        doc.EndUndo()
        # c4d.EventAdd()

    def IsSelected(self, root, userdata, obj):
        """
        Returns:
          (bool): True if *obj* is selected, False if not.
        """

        return obj.GetBit(c4d.BIT_ACTIVE)

    def DoubleClick(self, root, userdata, obj, col, mouseinfo):
        """
        Called when the user double-clicks on an entry in the TreeView.

        Returns:
          (bool): True if the double-click was handled, False if the
            default action should kick in. The default action will invoke
            the rename procedure for the object, causing `SetName()` to be
            called.
        """

        c4d.gui.MessageDialog("You clicked on " + obj.GetName())
        return True

    def IsResizeColAllowed(self, root, userdata, lColID):
        return True

    def IsTristate(self, root, userdata):
        return False

    def GetDragType(self, root, userdata, obj):
        """
        Returns:
          (int): The drag datatype.
        """

        return c4d.DRAGTYPE_ATOMARRAY

    def DragStart(self, root, userdata, obj):
        """
        Returns:
          (int): Bitmask specifying options for the drag, whether it is
            allowed, etc.
        """

        return c4d.TREEVIEW_DRAGSTART_ALLOW | c4d.TREEVIEW_DRAGSTART_SELECT

    def GetId(self, root, userdata, obj):
        """
        Return a unique ID for the element in the TreeView.
        """

        return obj.GetUniqueID()

    def SetName(self, root, userdata, obj, name):
        """
        Called when the user renames the element. `DoubleClick()` must return
        False for this to work.
        """

        obj.SetName(name)

    def DrawCell(self, root, userdata, obj, col, drawinfo, bgColor):
        """
        Called for a cell with layout type `c4d.LV_USER` or `c4d.LV_USERTREE`
        to draw the contents of a cell.
        """

        if col == ID_ICON:
            icon = obj.GetIcon()
            drawinfo["frame"].DrawBitmap(
                icon["bmp"], drawinfo["xpos"], drawinfo["ypos"],
                16, 16, icon["x"], icon["y"], icon["w"], icon["h"], c4d.BMP_ALLOWALPHA)

    def SetCheck(self, root, userdata, obj, column, checked, msg):
        """
        Called when the user clicks on a checkbox for an object in a
        `c4d.LV_CHECKBOX` column.
        """

        if checked:
            obj.SetBit(BIT_CHECKED)
        else:
            obj.DelBit(BIT_CHECKED)

        def checked_collect(op):
            if op is None:
                return
            elif isinstance(op, c4d.documents.BaseDocument):
                for obj in op.GetObjects():
                    for x in checked_collect(obj):
                        yield x
            else:
                if op.GetBit(BIT_CHECKED):
                    yield op
                for child in op.GetChildren():
                    for x in checked_collect(child):
                        yield x

        status = ', '.join(obj.GetName() for obj in checked_collect(obj.GetDocument()))
        self._dlg().SetString(1001, "Checked: " + status)
        self._dlg()._treegui.Refresh()

    def IsChecked(self, root, userdata, obj, column):
        """
        Returns:
          (int): Status of the checkbox in the specified *column* for *obj*.
        """

        val = obj.GetBit(BIT_CHECKED)
        if val == True:
            return c4d.LV_CHECKBOX_CHECKED | c4d.LV_CHECKBOX_ENABLED
        else:
            return c4d.LV_CHECKBOX_ENABLED

    def HeaderClick(self, root, userdata, column, channel, is_double_click):
        """
        Called when the TreeView header was clicked.

        Returns:
          (bool): True if the event was handled, False if not.
        """

        c4d.gui.MessageDialog("You clicked on the '%i' header." % (column))
        return True

    def AcceptDragObject(self, root, userdata, obj, dragtype, dragobject):
        """
        Called when a drag & drop operation hovers over the TreeView to check
        if the drag can be accepted.

        Returns:
          (int, bool)
        """

        if dragtype != c4d.DRAGTYPE_ATOMARRAY:
            return 0

        return c4d.INSERT_BEFORE | c4d.INSERT_AFTER | c4d.INSERT_UNDER, True

    def GenerateDragArray(self, root, userdata, obj):
        """
        Return:
          (list of c4d.BaseList2D): Generate a list of objects that can be
            dragged from the TreeView for the `c4d.DRAGTYPE_ATOMARRAY` type.
        """

        if obj.GetBit(c4d.BIT_ACTIVE):
            return [obj, ]

    def InsertObject(self, root, userdata, obj, dragtype, dragobject, insertmode, bCopy):
        """
        Called when a drag is dropped on the TreeView.
        """

        if dragtype != c4d.DRAGTYPE_ATOMARRAY:
            return  # Shouldnt happen, we catched that in AcceptDragObject

        for op in dragobject:
            op.Remove()
            if insertmode == c4d.INSERT_BEFORE:
                op.InsertBefore(obj)
            elif insertmode == c4d.INSERT_AFTER:
                op.InsertAfter(obj)
            elif insertmode == c4d.INSERT_UNDER:
                op.InsertUnder(obj)
        return

    def GetColumnWidth(self, root, userdata, obj, col, area):
        return 80  # All have the same initial width

    def IsMoveColAllowed(self, root, userdata, lColID):
        # The user is allowed to move all columns.
        # TREEVIEW_MOVE_COLUMN must be set in the container of AddCustomGui.
        return True

    def GetColors(self, root, userdata, obj, pNormal, pSelected):
        """
        Retrieve colors for the TreeView elements.

        Returns:
          (int or c4d.Vector, int or c4d.Vector): The colors for the normal
            and selected state of the element.
        """

        usecolor = obj[c4d.ID_BASEOBJECT_USECOLOR]
        if usecolor == c4d.ID_BASEOBJECT_USECOLOR_ALWAYS:
            pNormal = obj[c4d.ID_BASEOBJECT_COLOR]
        return pNormal, pSelected

    # Context menu entry for "Hello World!"
    ID_HELLOWORLD = 355435345

    def CreateContextMenu(self, root, userdata, obj, lColumn, bc):
        """
        User clicked with the right mouse button on an entry so
        we can here enhance the menu
        """

        bc.SetString(self.ID_HELLOWORLD, "Hello World!")

    def ContextMenuCall(self, root, userdata, obj, lColumn, lCommand):
        """
        The user executes an entry of the context menu.

        Returns:
          (bool): True if the event was handled, False if not.
        """

        if lCommand == self.ID_HELLOWORLD:
            obj.SetName("Hello World!")
            c4d.EventAdd()
            return True
        return False

    def SelectionChanged(self, root, userdata):
        """
        Called when the selected elements in the TreeView changed.
        """

        print "The selection changed"

    def Scrolled(self, root, userdata, h, v, x, y):
        """
        Called when the TreeView is scrolled.
        """

        self._dlg().SetString(1002, ("H: %i V: %i X: %i Y: %i" % (h, v, x, y)))


class TestDialog(c4d.gui.GeDialog):
    _treegui = None

    def CreateLayout(self):
        # Create the TreeView GUI.
        customgui = c4d.BaseContainer()
        customgui.SetBool(c4d.TREEVIEW_BORDER, True)
        customgui.SetBool(c4d.TREEVIEW_HAS_HEADER, True)
        customgui.SetBool(c4d.TREEVIEW_HIDE_LINES, False)
        customgui.SetBool(c4d.TREEVIEW_MOVE_COLUMN, True)
        customgui.SetBool(c4d.TREEVIEW_RESIZE_HEADER, True)
        customgui.SetBool(c4d.TREEVIEW_FIXED_LAYOUT, True)
        customgui.SetBool(c4d.TREEVIEW_ALTERNATE_BG, True)
        customgui.SetBool(c4d.TREEVIEW_CURSORKEYS, True)
        customgui.SetBool(c4d.TREEVIEW_NOENTERRENAME, False)

        self._treegui = self.AddCustomGui(
            1000, c4d.CUSTOMGUI_TREEVIEW, "", c4d.BFH_SCALEFIT | c4d.BFV_SCALEFIT,
            300, 300, customgui)
        if not self._treegui:
            print "[ERROR]: Could not create TreeView"
            return False

        self.AddMultiLineEditText(1001, c4d.BFH_SCALEFIT | c4d.BFV_SCALE | c4d.BFV_BOTTOM, 0, 40)
        self.AddEditText(1002, c4d.BFH_SCALEFIT, 0, 25)
        return True

    def CoreMessage(self, id, msg):
        if id == c4d.EVMSG_CHANGE:
            # Update the treeview on each change in the document.
            self._treegui.Refresh()
        return True

    def InitValues(self):
        # Initialize the column layout for the TreeView.
        layout = c4d.BaseContainer()
        layout.SetLong(ID_CHECKBOX, c4d.LV_CHECKBOX)
        layout.SetLong(ID_ICON, c4d.LV_USER)
        layout.SetLong(ID_TREEVIEW, c4d.LV_TREE)
        self._treegui.SetLayout(3, layout)

        # Set the header titles.
        self._treegui.SetHeaderText(ID_CHECKBOX, "Check")
        self._treegui.SetHeaderText(ID_ICON, "Icon")
        self._treegui.SetHeaderText(ID_TREEVIEW, "Name")
        self._treegui.Refresh()

        # Don't need to store the hierarchy, due to SetRoot()
        data_model = Hierarchy(self)
        self._treegui.SetRoot(None, data_model, None)
        self._treegui.SetFocusItem(None)

        self.SetString(1002, "Scroll values")
        return True


class MenuCommand(c4d.plugins.CommandData):
    dialog = None

    def Execute(self, doc):
        """
        Just create the dialog when the user clicked on the entry
        in the plugins menu to open it
        """

        if self.dialog is None:
            self.dialog = TestDialog()
        return self.dialog.Open(
            c4d.DLG_TYPE_ASYNC, PLUGIN_ID, defaulth=600, defaultw=600)

    def RestoreLayout(self, sec_ref):
        """
        Same for this method. Just allocate it when the dialog is needed.
        """

        if self.dialog is None:
            self.dialog = TestDialog()
        return self.dialog.Restore(PLUGIN_ID, secret=sec_ref)


def main():
    c4d.plugins.RegisterCommandPlugin(
        PLUGIN_ID, PLUGIN_NAME, 0, None, PLUGIN_HELP, MenuCommand())


if __name__ == "__main__":
    main()