# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""A tree widget with multiple selection and drag-and-drop support."""

# Derived from Python source code, which says:
# Copyright (c) 2001, 2002, 2003 Python Software Foundation.
# All rights reserved.

import os
from tkinter import dnd as Tkdnd
import tkinter as Tkinter

ICONDIR = "icons"

# Look for Icons subdirectory in the same directory as this module
try:
    _icondir = os.path.join(os.path.dirname(__file__), ICONDIR)
except NameError:
    _icondir = ICONDIR
if os.path.isdir(_icondir):
    ICONDIR = _icondir
elif not os.path.isdir(ICONDIR):
    raise RuntimeError("Can't find icon directory (%r)" % ICONDIR)


class TreeItem(object):

    """Abstract class representing tree items.

    Methods should typically be overridden, otherwise a default action
    is used.  There's no need to call the base class method, as the
    base class methods essentially do nothing interesting (except for
    get_icon()).

    XXX The drag-and-drop (DND) API is still experimental; it requires
    significant cooperation of some other part of the application.
    Exactly which part of the API should be defined here and which
    part should be defined elsewhere is unclear; I need to have played
    with this more.  Right now, DND only seems to work within the same
    widget anyway.

    """

    node = None  # XXX This creates cycles between nodes & items; do we care?

    def __init__(self):
        """Constructor."""

    def get_icon(self):
        """Returns the name of the icon to display.

        The default implementation returns a different icon when the
        node is selected, to emulate the Windows Explorer.

        """
        if self.node is not None and self.node.selected:
            return "openfolder"
        else:
            return "folder"

    def get_text(self):
        """Returns the text string to display."""
        return ""

    def is_expandable(self):
        """Returns whether there are subitems.

        This is separate from get_sublist() because sometimes
        get_sublist() can be expensive.  It's even possible that
        get_sublist() returns an empty list while is_expandable()
        returns True; but then the next call to is_expandable() should
        return False.

        """
        return False

    def get_sublist(self):
        """Returns the list of items forming the sublist.

        BEWARE: this *must* cache the items and return the same item
        for the same object each time, as long as the object is still
        considered a child; otherwise the widget tree will not display
        properly.

        """
        return []

    def on_selection(self):
        """Called when the item is added to the selection."""

    def on_deselection(self):
        """Called when the item is removed from the selection."""

    def on_action(self):
        """Called when the item is double-clicked or Return is hit."""

    def on_drop(self, source_items):
        """Called when dropping other items on this one."""

    def is_draggable(self):
        """Called to check whether this item can be dragged using drag-and-drop."""
        return False


class TreeNode(object):

    """Class for drawing a tree on a canvas.

    All attributes are considered read-only for public consumption.
    All modification of attributes must go through method calls.
    Their meaning (the first three are also the constructor
    arguments):

    - canvas: Tkinter.Canvas widget where the tree is drawn
    - parent: TreeNode giving this node's parent; None if this is the tree root
    - item: TreeItem subclass giving the node's associated item
    - expanded: bool indicating whether the node's subtree is visible
    - selected: bool indicating whether the node is part of the selection
    - children: list of TreeNodes giving this node's children; this
                may be non-empty even when expanded == False
    - x, y: ints giving the x and y position where this node is drawn;
            y may be None until the node is first drawn
    - root: TreeNode giving the root of the tree

    There are also internal attributes and methods; these have names
    starting with a single underscore.

    All public methods maintain the class invariant.  This invariant
    is hard to define precisely, but easy to describe informally: each
    node of the 'visible tree' is drawn at the correct location in the
    canvas with the correct plus/minus control, icon, text, and
    highlighting; and no other nodes are drawn.  The 'visible tree' is
    defined by starting at the root and descending only into children
    of *expanded* nodes.

    Whenever a node is expanded, its item is asked for its list of
    subitems (item.get_sublist()), and the node's list of subnodes
    (node.children) is recomputed from this; this process is applied
    recursively.  This means that if the contents of a subtree has
    been modified while the subtree was collapsed, the modifications
    are automatically revealed when it is expanded again.  But
    structural changes to the visible (expanded) part of the tree are
    not automatically recomputed; you must call node.refresh() to
    refresh a node's subtree.  As a convenience, node.expand() calls
    node.refresh() when the tree is already expanded.

    """

    # Style note: the methods in this class are grouped together by
    # functionality and presented roughly in a bottom-up order.

    # --- Constructor ---

    def __init__(self, canvas, parent, item):
        """Constructor.  See the class documentation for the arguments."""
        assert isinstance(canvas, Tkinter.Canvas) # XXX What about duck typing?
        self.canvas = canvas
        self.parent = parent
        self.item = item
        item.node = self
        self.expanded = False
        self.selected = False
        self.children = []
        if parent is None:
            self.x = self.y = 0
            self.root = self
            self._iconimages = {}
            self.refresh() # Ensure the invariant
        else:
            self.x = self.parent.x + 20
            self.y = None
            self.root = parent.root

    # --- Customization methods, intended to be overridden/extended --

    #@classmethod
    def factory(cls, canvas, parent, item):
        """Node factory class method.  Returns a new instance.

        The arguments are the same as for the constructor; see the
        class documentation for their documentation.

        This exists so that you can override it.

        """
        return cls(canvas, parent, item)
    factory = classmethod(factory) # XXX In Python 2.4, use a decorator instead

    def mouse_bindings(self):
        """Returns a list of (event, callback) pairs for mouse bindings.

        This exists so that you can override it.  Also because the
        bindings are referenced in two places.

        """
        return [
            ("<1>", self.on_click),
            ("<Shift-1>", self.on_shift_click),
            ("<Control-1>", self.on_control_click),
            ("<3>", self.on_click),
            ("<Double-1>", self.on_action),
            ]

    def keyboard_bindings(self):
        """Returns a list of (event, callback) pairs for keyboard bindings.

        This exists so that you can override it.

        """
        return [
            ("<Return>", self.on_action),
            ("<Left>", self.on_left),
            ("<Right>", self.on_right),
            ("<Up>", self.on_up),
            ("<Shift-Up>", self.on_shift_up),
            ("<Down>", self.on_down),
            ("<Shift-Down>", self.on_shift_down),
            ("<Home>", self.on_home),
            ("<Shift-Home>", self.on_shift_home),
            ("<End>", self.on_end),
            ("<Shift-End>", self.on_shift_end),
            ]

    # --- Accessor methods ---

    def __iter__(self):
        """Generator yielding all the nodes in this subtree.

        To clarify: this starts by yielding self; then recurses down
        the children but only if expanded is True.

        The consumer may modify the yielded node, affecting the
        iteration (see e.g. refresh()).

        """
        stack = [self]
        while stack:
            node = stack.pop()
            yield node
            if node.expanded:
                stack.extend(node.children[::-1])

    def successors(self):
        """Generator yielding all nodes below this node in the Y direction.

        To clarify: this node's children are excluded, but the
        children of later nodes are included.

        """
        q = self
        p = q.parent
        while p is not None:
            i = p.children.index(q) # May raise ValueError on tree corruption
            for sibling in p.children[i+1:]:
                for node in sibling:
                    yield node
            q = p
            p = q.parent

    def lastdescendant(self):
        """Returns this node's last descendant; maybe itself."""
        node = self
        while node.expanded and node.children:
            node = node.children[-1]
        return node

    def above(self):
        """Returns the node above this one (in the y dimension); or None."""
        parent = self.parent
        if parent is None:
            return None
        i = parent.children.index(self)
        if i == 0:
            return parent
        else:
            return parent.children[i-1].lastdescendant()

    def below(self):
        """Returns the node below this one (in the y dimension); or None."""
        if self.expanded and self.children:
            return self.children[0]
        for node in self.successors():
            return node
        return None

    def selection(self):
        """Returns a list containing the selected nodes in this subtree.

        This is not a generator since most uses want the concrete list.

        """
        return [node for node in self if node.selected]

    def find_node(self, x, y):
        """Returns the node at canvas coordinates (x, y); or None."""
        for node in self.root:
            if node.y <= y < node.y+17:
                return node
        return None

    # --- Drawing implementation methods; these are all internal ---

    def _geticonimage(self, name):
        """Internal: returns a named Tkinter.PhotoImage object.

        This maintains a cache in the  root node.

        """
        image = self.root._iconimages.get(name)
        if image is None:
            file, ext = os.path.splitext(name)
            ext = ext or ".gif"
            fullname = os.path.join(ICONDIR, file + ext)
            image = Tkinter.PhotoImage(master=self.canvas, file=fullname)
            self.root._iconimages[name] = image
        return image

    _plusid = None
    _plusname = None
    _plusbinding = None

    def _hideplus(self):
        """Internal: hides the '+' or '-' icon for expandable nodes."""
        self._plusname = None
        if self._plusid:
            if self._plusbinding:
                self.canvas.tag_unbind(self._plusid, "<1>", self._plusbinding)
                self._plusbinding = None
            id = self._plusid
            self._plusid = None
            self.canvas.delete(id)

    def _drawplus(self):
        """Internal: draws the '+' or '-' icon for expandable nodes."""
        if not self.item.is_expandable():
            self._hideplus()
        else:
            if self.expanded:
                iconname = "minusnode"
                callback = self.collapse
            else:
                iconname = "plusnode"
                callback = self.expand
            if self._plusid and self._plusbinding and iconname == self._plusname:
                return
            self._hideplus()
            image = self._geticonimage(iconname)
            id = self.canvas.create_image(self.x+9, self.y+8, image=image)
            self._plusid = id
            self._plusbinding = self.canvas.tag_bind(id, "<1>", callback)
            self._plusname = iconname

    _iconid = None
    _iconname = None
    _iconbindings = None

    def _hideicon(self):
        """Internal: hides the icon."""
        self._iconname = None
        if self._iconid:
            if self._iconbindings:
                for event, funcid in self._iconbindings:
                    self.canvas.tag_unbind(self._iconid, event, funcid)
                self._iconbindings = None
            id = self._iconid
            self._iconid = id
            self.canvas.delete(id)

    def _drawicon(self):
        """Internal: draws the icon."""
        iconname = self.item.get_icon()
        if self._iconid and self._iconbindings and iconname == self._iconname:
            return
        self._hideicon()
        image = self._geticonimage(iconname)
        id = self.canvas.create_image(self.x+20, self.y+1, anchor="nw", image=image)
        self._iconid = id
        iconbindings = []
        for event, callback in self.mouse_bindings():
            iconbindings.append((event, self.canvas.tag_bind(id, event, callback)))
        self._iconbindings = iconbindings
        self._iconname = iconname

    _textid = None
    _textwidget = None

    def _hidetext(self):
        """Internal: hides the text."""
        if self._textid:
            id = self._textid
            self._textid = None
            self.canvas.delete(id)
            self._textid = None
        if self._textwidget:
            widget = self._textwidget
            self._textwidget = None
            widget.destroy()

    def _drawtext(self):
        """Internal: draws the text."""
        widget = self._textwidget
        if widget is None:
            widget = Tkinter.Label(self.canvas, bd=0, padx=2, pady=2)
            self._bg = self.canvas["bg"]
            self._fg = widget["fg"]
            self._textwidget = widget
            for event, callback in self.mouse_bindings() + self.keyboard_bindings():
                widget.bind(event, callback)
        text = self.item.get_text()
        if self.selected:
            fg = self.canvas["selectforeground"]
            bg = self.canvas["selectbackground"]
            takeFocus = True
        else:
            fg = self._fg
            bg = self._bg
            takeFocus = False
        widget.configure(fg=fg, bg=bg, text=text, takefocus=takeFocus)
        if takeFocus:
            widget.focus_set()
        if not self._textid:
            self._textid = self.canvas.create_window(self.x+39, self.y,
                                                     anchor="nw", window=widget)

    def _hide(self):
        """Internal: hide this node."""
        self._hideplus()
        self._hideicon()
        self._hidetext()
        self.y = None

    def _hideall(self, keep=None):
        """Hides this node and all its descendants.

        The optional argument specifies a node that should remain
        visible.

        """
        for node in self:
            if node is not keep:
                node._hide()

    def _draw(self, y=None):
        """Redraws this node in place.

        If the y argument is given (not None), it specifies a new y
        position for the node.

        Returns the y coordinate for the *next* node to be drawn.

        """
        if y is not None:
            if self.y is not None and y is not self.y:
                self._moveto(y)
            self.y = y
        assert self.y is not None
        self._drawplus()
        self._drawicon()
        self._drawtext()
        return self.y+17

    def _moveto(self, y):
        """Internal: moves this node to the given y position."""
        assert self.y is not None
        dy = y - self.y
        if dy:
            for id in self._plusid, self._iconid, self._textid:
                if id:
                    self.canvas.move(id, 0, dy)
        self.y = y
        return y+17

    def _reload(self):
        """Internal: recomputes the set of children of an expanded node."""
        if self.expanded:
            old = dict([(node.item, node) for node in self.children])
            self.children = []
            for item in self.item.get_sublist():
                node = old.get(item)
                if node is None:
                    node = self.factory(self.canvas, self, item)
                else:
                    del old[item]
                self.children.append(node)
            for node in old.values():
                node._hideall()
            if not self.children:
                self.expanded = False

    def refresh(self):
        """Redraws this node and its children, e.g. after expand/collapse.

        This recomputes the descendants.

        The successors are moved up/down by the appropriate amount,
        too (but otherwise assumed to be unchanged).

        Finally, the scroll region is updated.

        """
        y = self.y
        for node in self:
            node._reload()
            y = node._draw(y)
        for node in self.successors():
            if node.y == y:
                break
            y = node._moveto(y)
        x0, y0, x1, y1 = self.canvas.bbox(Tkinter.ALL)
        self.canvas.configure(scrollregion=(0, 0, x1, y1))

    def view(self):
        """Adjusts the view to ensuer this node is visible.

        This attempts to show 10 pixels towards the edge; that gives
        the user a hint on whether there's more there.

        Much pain and effort went into finding the right algorithm to
        decide when to scroll and to where.

        """
        EXTRA = 10 # Extra space we like to keep towards the edge
        scrollregion = self.canvas["scrollregion"].split()
        y1 = float(scrollregion[-1])
        ytop, ybot = self.canvas.yview()
        ytop *= y1
        ybot *= y1
        if self.y >= ytop and self.y+17 <= ybot:
            return
        scrollregion = self.canvas["scrollregion"].split()
        y1 = float(scrollregion[-1])
        if self.y < ytop:
            y = self.y - EXTRA
        else:
            y = self.y + 17 + EXTRA - (ybot - ytop)
        y = max(0, y)
        self.canvas.yview_moveto(y / y1)

    def focus(self):
        """Sets the keyboard focus on this node, if possible."""
        if self._textwidget:
            self._textwidget.focus_set()

    # --- State modifying methods ---

    # These generally operate on the current subtree, but some affect
    # the whole tree (select

    def expand(self, event=None):
        # Takes an event because it's a callback for the plus icon
        """Expands this node.

        This restores the entire subtree that was visible when the
        node was last collapsed.

        If the item was already expanded, we still do the
        recomputation; this seems to be the most common usage pattern.

        """
        if not self.expanded and self.item.is_expandable():
            self.expanded = True
        self.refresh()

    def collapse(self, event=None):
        # Takes an event because it's a callback for the minus icon
        """Collapses this node.

        The subtree is hidden but its state remembered until it is
        next expanded.

        """
        if self.expanded:
            self.deselectall(self)
            self._hideall(self)
            self.expanded = False
            self.refresh()

    def selectadd(self):
        """Adds this node to the selection."""
        if not self.selected:
            self.selected = True
            self._draw()
            self.item.on_selection()
        self.focus()
        self.view()

    def deselect(self):
        """Removes this node from the selection."""
        if self.selected:
            self.selected = False
            self._draw()
            self.item.on_deselection()

    def deselectall(self, keep=None):
        """Removes this node and all its descendants from the selection.

        The optional argument specifies a node that should remain
        selected.

        """
        for node in self.selection():
            if node is not keep:
                node.deselect()

    def select(self):
        """Sets the selection to just this node."""
        self.root.deselectall(self)
        self.selectadd()

    def selectrange(self, other=None):
        """Sets the selection to a continuous range.

        One end of the range is the current node.  The other end is an
        optional argument, defaulting to the end of the selection
        that's farthest away.

        """
        # There are a zillion ways to define shift-click.  My
        # definition: reset the selection to all nodes between the
        # clicked node and the end of the selection that is farthest
        # away from the clicked node, inclusive the end points. This
        # works for simple cases, and, frankly, that's all I care
        # about, given what a hairy mess this is and the fact that no
        # two applications seem to do this the same anyway.  If
        # there's *no* current selection, simply ignore the shift key.
        if other is None:
            selection = self.root.selection()
            if len(selection) == 0:
                self.select()
                return
            top = selection[0]
            bot = selection[-1]
            dtop = abs(self.y - top.y)
            dbot = abs(self.y - bot.y)
            if dtop >= dbot:
                bot = self
            else:
                top = self
        else:
            top, bot = self, other
        if top.y > bot.y:
            top, bot = bot, top
        for node in self.root:
            if node.y < top.y or node.y > bot.y:
                node.deselect()
            elif not node.selected:
                node.selectadd()
        self.selectadd()

    # --- Methods implementing mouse event bindings ---

    def on_click(self, event=None):
        if not self.selected:
            self.select()
        for node in self.root.selection():
            if not node.item.is_draggable():
                return
        Tkdnd.dnd_start(self, event)

    def on_shift_click(self, event=None):
        self.selectrange()

    def on_control_click(self, event=None):
        if self.selected:
            self.deselect()
        else:
            self.selectadd()

    def on_action(self, event=None):
        self.item.on_action()

    # --- Methods implementing keyboard event bindings ---

    def on_left(self, event=None):
        if self.expanded:
            self.collapse()
        elif self.parent is not None:
            self.parent.select()

    def on_right(self, event=None):
        if not self.expanded:
            self.expand()
        elif self.children:
            self.children[0].select()

    def on_up(self, event=None):
        selection = self.root.selection()
        if len(selection) > 1:
            selection[0].select()
            return
        node = self.above()
        if node is not None:
            node.select()

    def on_shift_up(self, event=None):
        node = self.above()
        if node is not None:
            node.selectrange()

    def on_down(self, event=None):
        selection = self.root.selection()
        if len(selection) > 1:
            selection[-1].select()
            return
        node = self.below()
        if node is not None:
            node.select()

    def on_shift_down(self, event=None):
        node = self.below()
        if node is not None:
            node.selectrange()

    def on_home(self, event=None):
        self.root.select()

    def on_shift_home(self, event=None):
        self.selectrange(self.root)

    def on_end(self, event=None):
        self.root.lastdescendant().select()

    def on_shift_end(self, event=None):
        self.selectrange(self.root.lastdescendant())

    # --- Drag-and-drop support methods ---

    def dnd_end(self, target, event):
        self.dnd_leave(None, event)
        if target is not None:
            selection = self.root.selection()
            if selection and target not in selection:
                source_items = [node.item for node in selection]
                target.item.on_drop(source_items, event)
            return

    def dnd_enter(self, source, event):
        self._textwidget.configure(fg="white", bg="#0000B0")

    def dnd_motion(self, source, event):
        pass

    def dnd_leave(self, source, event):
        self._drawtext()

    def dnd_commit(self, source, event):
        self.dnd_leave(source, event)


class Traversal(object):

    """Helper class to maintain state for item traversal."""

    def __init__(self, item):
        self.cache = {}
        if item is None:
            self.path = []
        else:
            path = [item]
            while item.node is not None:
                parentnode = item.node.parent
                if parentnode is None:
                    break
                item = parentnode.item
                if item is None:
                    break
                path.append(item)
            path.reverse()
            self.path = path

    def next(self):
        """Called to find the next item in a pre-order traversal.

        Raises ValueError if this node's item is not in its parent's
        sublist.

        """
        if not self.path:
            return None
        item = self.path[-1]
        if item.is_expandable():
            sublist = self.cache.get(item)
            if sublist is None:
                sublist = item.get_sublist() or []
                self.cache[item] = sublist
            if sublist:
                item = sublist[0]
                self.path.append(item)
                return item
        while self.path:
            item = self.path.pop()
            if not self.path:
                return None
            parent = self.path[-1]
            sublist = self.cache.get(parent)
            if sublist is None:
                sublist = parent.get_sublist() or []
                self.cache[parent] = sublist
            i = sublist.index(item) + 1
            if i < len(sublist):
                item = sublist[i]
                self.path.append(item)
                return item


def test():
    # A mini directory browser.
    root = Tkinter.Tk()
    root.wm_title(os.getcwd())
    canvas = Tkinter.Canvas(root, relief="sunken", borderwidth=2,
                            highlightthickness=2)
    canvas.pack(expand=True, fill="both")
    class FSTreeItem(TreeItem):
        def __init__(self, fn):
            self.fn = fn
            self.cache = {}
            TreeItem.__init__(self)
        def get_text(self):
            return os.path.basename(self.fn)
        def get_icon(self):
            if os.path.isdir(self.fn):
                return TreeItem.get_icon(self)
            else:
                return "python"
        def is_expandable(self):
            return os.path.isdir(self.fn) and bool(os.listdir(self.fn))
        def get_sublist(self):
            names = os.listdir(self.fn)
            names.sort(key=str.lower)
            lst = []
            for name in names:
                item = self.cache.get(name)
                if item is None:
                    item = FSTreeItem(os.path.join(self.fn, name))
                    self.cache[name] = item
                lst.append(item)
            return lst
    topnode = TreeNode(canvas, None, FSTreeItem(os.getcwd()))
    topnode.select()
    root.mainloop()


if __name__ == "__main__":
    test()
