TreeView C4D Python [Material] v1.2

"""Example for setting the column width of a tree view by the width of its
items.

The only changes which have been made in comparison to the original code can
be found in TextureObject.__init__(), TextureObject.RandomString() and
ListView.GetColumnWidth(). The tree view items #TextureObject now generate
random #texturePath and #otherData strings and the treeview scales its columns
to the length of these strings.

Adapted from:
    https://plugincafe.maxon.net/topic/10654/14102_using-customgui-listview/2
"""
import c4d
import random
import sys

# Be sure to use a unique ID obtained from http://www.plugincafe.com/.
PLUGIN_ID = 1000010 # TEST ID ONLY

# TreeView Column IDs.
ID_CHECKBOX = 1
ID_NAME = 2
ID_OTHER = 3
ID_LONGFILENAME = 4

# A tuple of characters to select from (a-z).
CHARACTERS = tuple( chr ( n) for n in range(97, 122))
doc = c4d.documents.GetActiveDocument()

class MaterialCollection:
    """Handles a collection of materials from a singular document for long-term referencing.

    Will reestablish material references when possible.
    """

    def __init__(self, collection: list[c4d.BaseMaterial]):
        """Initializes the collection with a list of materials.
        """
        if not isinstance(collection, list):
            raise TypeError(f"{collection = }")

        # The list of materials, the list of markers for them, and their associated document.
        self._materials: list[c4d.BaseMaterial] = []
        self._markers: list[bytes] = []
        self._doc: c4d.documents.BaseDocument = None

        for material in collection:
            # Sort out illegal types, "dead" materials, and materials not attached to a document.
            if not isinstance(material, c4d.BaseMaterial) or not material.IsAlive():
                raise TypeError(f"{material} is not alive or not a BaseMaterial.")
            if material.GetDocument() is None:
                raise RuntimeError("Dangling material cannot be managed.")

            # Establish what the document of the materials is and that they all belong to the same
            # document. Technically, a BaseDocument can also be dead (when it has been closed),
            # but this is less likely to happen, so a stored reference to it can also become dead.
            if self._doc is None:
                self._doc = material.GetDocument()
            elif self._doc != material.GetDocument():
                raise RuntimeError("Can only manage materials attached to the same document.")

            # Append the material and its UUID marker to the internal lists.
            self._materials.append(material)
            self._markers.append(MaterialCollection.GetUUID(material))

    def __len__(self):
        """Returns the number of managed materials.
        """
        return len(self._materials)

    def __getitem__(self, index: int) -> c4d.BaseMaterial:
        """Returns the material at #index.
        """
        if len(self._materials) - 1 < index:
            raise IndexError(f"The index '{index}' is out of bounds.")

        # The item is still alive, we can return it directly.
        item: c4d.BaseMaterial = self._materials[index]
        if item.IsAlive():
            print ("Returning still valid material reference")
            return item

        # The item is not alive anymore, try to find its new instance.
        if not self._doc.IsAlive():
            raise RuntimeError(
                f"{self._doc} is not valid anymore, cannot reestablish material references.")

        # Get the marker identifying the material at #index.
        materialMarker: bytes = self._markers[index]

        # Iterate over all materials in #_doc to find one with a matching marker.
        for material in self._doc.GetMaterials():
            otherMarker: bytes = MaterialCollection.GetUUID(material)

            # We found a match, update the internal data, and return the material.
            if materialMarker == otherMarker:
                self._materials[index] = material
                print ("Returning re-established material reference")
                return material

        # There is no material with this marker anymore.
        raise RuntimeError(
            f"The material at {index} has not valid anymore and cannot be re-established.")

    def __iter__(self):
        """Iterates over all materials in the collection.
        """
        for index in range(len(self._materials)):
            yield self[index]

    @staticmethod
    def GetUUID(atom: c4d.C4DAtom) -> bytes:
        """Retrieves the UUID of an atom as a bytes object.

        Args:
            item (c4d.C4DAtom): The atom to get the UUID hash for.

        Raises:
            RuntimeError: On type assertion failure.

        Returns:
            bytes: The UUID of #atom.
        """
        if not isinstance(atom, c4d.C4DAtom):
            raise TypeError(f"{atom = }")

        # Get the MAXON_CREATOR_ID marker uniquely identifying the atom.
        uuid: memoryview = atom.FindUniqueID(c4d.MAXON_CREATOR_ID)
        if not isinstance(uuid, memoryview):
            raise RuntimeError(f"Illegal non-marked atom: {atom}")

        return bytes(uuid)

class TextureObject(object):
    """
    Class which represent a texture, aka an Item in our list
    """
    texturePath = "TexPath"
    otherData = "OtherData"
    _selected = False

    def __init__(self, texturePath):
        self.texturePath = texturePath
        self.otherData += texturePath

    @property
    def IsSelected(self):
        return self._selected

    def Select(self):
        self._selected = True

    def Deselect(self):
        self._selected = False

    def __repr__(self):
        return str(self)

    def __str__(self):
        return self.texturePath


class ListView(c4d.gui.TreeViewFunctions):


    def __init__(self, collection):

        self.listOfTexture = collection

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

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

    def GetColumnWidth(self, root, userdata, obj, col, area):
        """Measures the width of cells.

        Although this function is called #GetColumnWidth and has a #col, it is
        not only executed by column but by cell. So, when there is a column
        with items requiring the width 5, 10, and 15, then there is no need
        for evaluating all items. Each item can return its ideal width and
        Cinema 4D will then pick the largest value.

        Args:
            root (any): The root node of the tree view.
            userdata (any): The user data of the tree view.
            obj (any): The item for the current cell.
            col (int): The index of the column #obj is contained in.
            area (GeUserArea): An already initialized GeUserArea to measure
             the width of strings.

        Returns:
            TYPE: Description
        """
        # The default width of a column is 80 units.
        width = 80
        # Replace the width with the text width. area is a prepopulated
        # user area which has already setup all the font stuff, we can
        # measure right away.

        if col == ID_NAME:
            return area.DrawGetTextWidth(obj.texturePath) + 5
        if col == ID_OTHER:
            return area.DrawGetTextWidth(obj.otherData) + 5
        if col == ID_LONGFILENAME:
            return area.DrawGetTextWidth(obj.longfilename) + 5

        return 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 GetFirst(self, root, userdata):
        """
        Return the first element in the hierarchy, or None if there is no element.
        """
        rValue = None if not self.listOfTexture else self.listOfTexture[0]
        return rValue

    def GetDown(self, root, userdata, obj):
        """
        Return a child of a node, since we only want a list, we return None everytime
        """
        return None

    def GetNext(self, root, userdata, obj):
        """
        Returns the next Object to display after arg:'obj'
        """
        rValue = None
        currentObjIndex = self.listOfTexture.index(obj)
        nextIndex = currentObjIndex + 1
        if nextIndex < len(self.listOfTexture):
            rValue = self.listOfTexture[nextIndex]

        return rValue

    def GetPred(self, root, userdata, obj):
        """
        Returns the previous Object to display before arg:'obj'
        """
        rValue = None
        currentObjIndex = self.listOfTexture.index(obj)
        predIndex = currentObjIndex - 1
        if 0 <= predIndex < len(self.listOfTexture):
            rValue = self.listOfTexture[predIndex]

        return rValue

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

    def Select(self, root, userdata, obj, mode):
        """
        Called when the user selects an element.
        """
        if mode == c4d.SELECTION_NEW:
            for tex in self.listOfTexture:
                tex.Deselect()
            obj.Select()
        elif mode == c4d.SELECTION_ADD:
            obj.Select()
        elif mode == c4d.SELECTION_SUB:
            obj.Deselect()

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

    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.Select()
        else:
            obj.Deselect()

    def IsChecked(self, root, userdata, obj, column):
        """
        Returns: (int): Status of the checkbox in the specified *column* for *obj*.
        """
        if obj.IsSelected:
            return c4d.LV_CHECKBOX_CHECKED | c4d.LV_CHECKBOX_ENABLED
        else:
            return c4d.LV_CHECKBOX_ENABLED

    def GetName(self, root, userdata, obj):
        """
        Returns the name to display for arg:'obj', only called for column of type LV_TREE
        """
        return str(obj)  # Or obj.texturePath

    def DrawCell(self, root, userdata, obj, col, drawinfo, bgColor):
        """
        Draw into a Cell, only called for column of type LV_USER
        """
        if col in (ID_OTHER, ID_LONGFILENAME):
            text = obj.otherData if col == ID_OTHER else obj.longfilename
            canvas = drawinfo["frame"]
            textWidth = canvas.DrawGetTextWidth(text)
            textHeight = canvas.DrawGetFontHeight()
            xpos = drawinfo["xpos"]
            ypos = drawinfo["ypos"] + drawinfo["height"]

            if (drawinfo["width"] < textWidth):
                while (drawinfo["width"] < textWidth):
                    if len(text) <= 4:
                        text = "..."
                        break
                    text = text[:-4] + "..."
                    textWidth = canvas.DrawGetTextWidth(text)

            textWidth = canvas.DrawGetTextWidth(text)
            drawinfo["frame"].DrawText(text, xpos, ypos - int(textHeight * 1.1))

    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 " + str(obj))
        return True

    def DeletePressed(self, root, userdata):
        "Called when a delete event is received."
        for tex in reversed(self.listOfTexture):
            if tex.IsSelected:
                self.listOfTexture.remove(tex)

class TestDialog(c4d.gui.GeDialog):
    _treegui = None  # Our CustomGui TreeView

    collection: MaterialCollection = getattr(sys, "matCollection", None)
    if collection is None and doc.GetMaterials() != []:
        collection = MaterialCollection(doc.GetMaterials())
        setattr(sys, "matCollection", collection)
        print ("Stored material collection at sys.matCollection")
    elif collection is not None:
        print ("Retrieved stored material collection from sys.matCollection")
    else:
        print ("Document contains no materials.")
    _listView = ListView(collection)  # Our Instance of c4d.gui.TreeViewFunctions
    _listView._doc = c4d.documents.GetActiveDocument()

    def CreateLayout(self):
        # Create the TreeView GUI.
        customgui = c4d.BaseContainer()
        customgui.SetBool(c4d.TREEVIEW_BORDER, c4d.BORDER_THIN_IN)
        customgui.SetBool(c4d.TREEVIEW_HAS_HEADER, True)  # True if the tree view may have a header line.
        customgui.SetBool(c4d.TREEVIEW_HIDE_LINES, True)  # True if no lines should be drawn.
        customgui.SetBool(c4d.TREEVIEW_MOVE_COLUMN, True)  # True if the user can move the columns.
        customgui.SetBool(c4d.TREEVIEW_RESIZE_HEADER, True)  # True if the column width can be changed by the user.
        customgui.SetBool(c4d.TREEVIEW_FIXED_LAYOUT, True)  # True if all lines have the same height.
        customgui.SetBool(c4d.TREEVIEW_ALTERNATE_BG, True)  # Alternate background per line.
        customgui.SetBool(c4d.TREEVIEW_CURSORKEYS, True)  # True if cursor keys should be processed.
        customgui.SetBool(c4d.TREEVIEW_NOENTERRENAME, False)  # Suppresses the rename popup when the user presses enter.

        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.AddButton(1001, c4d.BFH_CENTER, name="Add")
        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_NAME, c4d.LV_TREE)
        layout.SetLong(ID_LONGFILENAME, c4d.LV_USER)
        layout.SetLong(ID_OTHER, c4d.LV_USER)
        self._layout = layout
        self._treegui.SetLayout(4, layout)

        # Set the header titles.
        self._treegui.SetHeaderText(ID_CHECKBOX, "Check")
        self._treegui.SetHeaderText(ID_NAME, "Name")
        self._treegui.SetHeaderText(ID_LONGFILENAME, "Long Filename")
        self._treegui.SetHeaderText(ID_OTHER, "Other")

        self._treegui.Refresh()

        # Set TreeViewFunctions instance used by our CUSTOMGUI_TREEVIEW

        self._treegui.SetRoot(self._treegui, self._listView, None)
        return True

    def Command(self, id, msg):
        # Click on button
        if id == 1001:

            '''
            # Add data to our DataStructure (ListView)
            newID = int(len(self._listView.listOfTexture) + 1)
            # print(newID)
            # newID = "T{}".format(newID)
            # print(newID)
            # tex = TextureObject(newID)
            tex = TextureObject()
            tex.texturePath = "Some new data"
            tex.longfilename = TextureObject.RandomString(20, 40)
            self._listView.listOfTexture.append(tex)
            self._treegui.Refresh()
            '''
            pass

        return True

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

    def Execute(self, doc):
        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):
        if self.dialog is None:
            self.dialog = TestDialog()
        return self.dialog.Restore(PLUGIN_ID, secret=sec_ref)

def main():
    c4d.plugins.RegisterCommandPlugin(
        PLUGIN_ID, "Python TreeView Example REVISED", 0, None, "Python TreeView Example", MenuCommand())


if __name__ == "__main__":
    main()