binarytree.py

import functools
from reprlib import recursive_repr
from typing import Iterable


class TreeOverflow(OverflowError, ValueError):
    pass


def _solve_args(**kwds):
    args = []
    left, right, left2, right2 = (kwds.pop(n, None) for n in ('left', 'right', 'l', 'r'))
    for long, short in ((left, left2), (right, right2)):
        if None in (long, short):
            args.append(long or short)
        else:
            short_n = (long_n := ("right", "left")[not args])[0]
            raise ValueError(f"can't specify both {long_n} and {short_n}")
    if args[0] is None and args[-1] is not None:
        args.reverse()
    elif isinstance(args[-1], Tree) and not isinstance(args[0], Tree):
        args.reverse()
    if kwds:
        raise ValueError(
            f'invalid argument{"s" if len(kwds) > 1 else ""}: {", ".join(kwds)}'
        )
    return args


@functools.total_ordering
class Tree:
    def __init__(self, left=None, right=None, **kwargs):
        self.__left = None
        self.__right = None
        self._depth = 0
        self.__walked = True
        self(left=left, right=right, **kwargs)

    @property
    def depth(self):
        if not self.__walked:
            self._walk(list)
        return self._depth

    @property
    def left(self):
        return self.__left

    @left.setter
    def left(self, value):
        self.__left = value

    @property
    def right(self):
        return self.__right

    @right.setter
    def right(self, value):
        self.__right = value

    def add_node(
            self,
            left=None,
            right=None,
            coerce_insertion=False,
            append_on_overflow=True,
            **kwargs
    ):
        left, right = args = _solve_args(left=left, right=right, **kwargs)
        if None not in args and (None, None) == (
            self.to_tuple(False) if not coerce_insertion else (None, None)
        ):
            self.left, self.right = sorted(args) if type(left) == type(right) else args
        elif (node := left or right) is not None:
            if self.left is None:
                if self.right is not None and self.right < node:
                    return self(left=self.right, right=node, coerce_insertion=True)
                self.left = node
            elif self.right is None:
                if self.left > node:
                    return self(left=node, right=self.left, coerce_insertion=True)
                self.right = node
            elif isinstance(self.left, Tree) and append_on_overflow:
                self.left(
                    node,
                    coerce_insertion=coerce_insertion,
                    append_on_overflow=append_on_overflow
                )
            elif isinstance(self.right, Tree) and append_on_overflow:
                self.right(
                    node,
                    coerce_insertion=coerce_insertion,
                    append_on_overflow=append_on_overflow
                )
            else:
                raise TreeOverflow('there is no place for this node')
        self.__walked = False
        return self

    __call__ = add_node

    def to_list(self, recursive=True):
        return self._walk(list, recursive)

    def to_tuple(self, recursive=True):
        return self._walk(tuple, recursive)

    def _walk(self, fn, recursive=True):
        res = fn(self._iwalk(fn, recursive))
        self._depth -= 1
        self.__walked = True
        return res

    def _iwalk(self, fn, recursive=True, _depth=0):
        for n in (self.left, self.right):
            if isinstance(n, (Iterable, Tree)) and n is not self and recursive:
                if isinstance(n, Tree):
                    _depth += 1
                    yield fn(n._iwalk(fn, _depth=_depth))
                    self._depth += n.depth
                else:
                    yield fn(n)
            else:
                print(' ' * _depth + str(n))
                yield n
        self._depth += 1
        self.__walked = True

    def __lt__(self, other):
        if isinstance(other, type(self)):
            other = other.to_tuple()
        elif not isinstance(other, Iterable):
            return True
        return self.to_tuple() < other

    def __eq__(self, other):
        if isinstance(other, type(self)):
            other = other.to_tuple()
        elif not isinstance(other, Iterable):
            return False
        return self.to_tuple() == other

    @recursive_repr('this')
    def __repr__(self):
        body = ", ".join(
            f"{k}={v}" for k, v in (('l', self.left), ('r', self.right)) if v is not None
        )
        return 'T' + body.join('()')


T = Tree