treeBase.py

"""Tree objects that support the `with` statement.
A bit of an OOP mess though.
"""
import contextvars
defaultParent = contextvars.ContextVar('defaultParent', default=None)
treeSessionStack=[]

from collections import UserList
import weakref

class TreeNode:
    def __init__(self,parent=None):
        if parent != None:
            self.parent=parent
        elif defaultParent:
            self.parent=defaultParent.get()
        else:
            self.parent=None
        super().__init__()

    @property
    def parent(self):
        return self._parent

    @parent.setter
    def parent(self,obj):
        self._reparent(obj,append=True)

    def _reparent(self,obj,append=True):
        """This bit of code is responsible for ensuring that TreeNodes have 1, and only 1,
        parent object.
        """
        #ToDo, the control flow on this is awful and dumb. Fix it.

        #Need to handle oldParent cleanup at the end or else the object de-references itself and
        # gets garbage collected... Or at least I think that's what's going on.
        oldParent=getattr(self,"_parent",None)

        # Make sure object is a supported type
        try: assert isinstance(obj,TreeContainer) or obj == None
        except:
            raise ValueError(f"Leaf nodes can only have TreeContainers as their parent objects, not {type(obj)}'s like {obj}")

        #Add object to new container and set backreference
        if append and obj != None:
            obj.append(self)

        if obj != None:
            #We use weakref to avoid reference loops
            self._parent = weakref.proxy(obj)
        else:
            self._parent = None

        #Remove item from old container
        if oldParent:
            while self in oldParent:
                oldParent.remove(self)

import contextlib

class TreeContainer(TreeNode,UserList):
    def __init__(self,data=list(),*args,**kwargs):
        super().__init__(*args,**kwargs)
        self.data.extend(data)
        self._session_token=None

    @contextlib.contextmanager
    def __call__(self):
        #This version is re-entrant, but needs to be called explicitly...
        # that could probably be hacked around, but it would be even more ugly.
        session_token=defaultParent.set(self)
        yield self
        defaultParent.reset(session_token)

    def append(self, value):
        super().append(value)
        try: assert isinstance(value,TreeNode)
        except: raise ValueError("TreeContainers can only contain TreeNodes")
        value._reparent(self,append=False)

t = TreeContainer()
with t():
    t2 = TreeContainer(['t2',],)
    TreeContainer(['t2-sub',],parent=t2)

    t3 = TreeContainer(['t3',],)
    with t3():
        TreeContainer(['t3-sub',],)

    t4 = TreeContainer(['t4',],)
    t4.append(TreeContainer(['t4-sub',],))

    t5 = TreeContainer(['t5',],)
    tTemp = TreeContainer(['t5-sub',],)
    tTemp.parent=t5

print(t)

#It's almost a test case!
assert repr(t) == "[['t2', ['t2-sub']], ['t3', ['t3-sub']], ['t4', ['t4-sub']], ['t5', ['t5-sub']]]"