treenode.py

from __future__ import annotations

import sys
from collections.abc import Iterator, Mapping
from pathlib import PurePosixPath
from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    TypeVar,
)

from xarray.core.utils import Frozen, is_dict_like

if TYPE_CHECKING:
    from xarray.core.types import T_DataArray


class InvalidTreeError(Exception):
    """Raised when user attempts to create an invalid tree in some way."""


class NotFoundInTreeError(ValueError):
    """Raised when operation can't be completed because one node is not part of the expected tree."""


class NodePath(PurePosixPath):
    """Represents a path from one node to another within a tree."""

    def __init__(self, *pathsegments):
        if sys.version_info >= (3, 12):
            super().__init__(*pathsegments)
        else:
            super().__new__(PurePosixPath, *pathsegments)
        if self.drive:
            raise ValueError("NodePaths cannot have drives")

        if self.root not in ["/", ""]:
            raise ValueError(
                'Root of NodePath can only be either "/" or "", with "" meaning the path is relative.'
            )
        # TODO should we also forbid suffixes to avoid node names with dots in them?


Tree = TypeVar("Tree", bound="TreeNode")


class TreeNode(Generic[Tree]):
    """
    Base class representing a node of a tree, with methods for traversing and altering the tree.

    This class stores no data, it has only parents and children attributes, and various methods.

    Stores child nodes in an dict, ensuring that equality checks between trees
    and order of child nodes is preserved (since python 3.7).

    Nodes themselves are intrinsically unnamed (do not possess a ._name attribute), but if the node has a parent you can
    find the key it is stored under via the .name property.

    The .parent attribute is read-only: to replace the parent using public API you must set this node as the child of a
    new parent using `new_parent.children[name] = child_node`, or to instead detach from the current parent use
    `child_node.orphan()`.

    This class is intended to be subclassed by DataTree, which will overwrite some of the inherited behaviour,
    in particular to make names an inherent attribute, and allow setting parents directly. The intention is to mirror
    the class structure of xarray.Variable & xarray.DataArray, where Variable is unnamed but DataArray is (optionally)
    named.

    Also allows access to any other node in the tree via unix-like paths, including upwards referencing via '../'.

    (This class is heavily inspired by the anytree library's NodeMixin class.)

    """

    _parent: Tree | None
    _children: dict[str, Tree]

    def __init__(self, children: Mapping[str, Tree] | None = None):
        """Create a parentless node."""
        self._parent = None
        self._children = {}

        if children:
            # shallow copy to avoid modifying arguments in-place (see GH issue #9196)
            self.children = {name: child.copy() for name, child in children.items()}

    @property
    def parent(self) -> Tree | None:
        """Parent of this node."""
        return self._parent

    @parent.setter
    def parent(self: Tree, new_parent: Tree) -> None:
        raise AttributeError(
            "Cannot set parent attribute directly, you must modify the children of the other node instead using dict-like syntax"
        )

    def _set_parent(
        self, new_parent: Tree | None, child_name: str | None = None
    ) -> None:
        # TODO is it possible to refactor in a way that removes this private method?

        if new_parent is not None and not isinstance(new_parent, TreeNode):
            raise TypeError(
                "Parent nodes must be of type DataTree or None, "
                f"not type {type(new_parent)}"
            )

        old_parent = self._parent
        if new_parent is not old_parent:
            self._check_loop(new_parent)
            self._detach(old_parent)
            self._attach(new_parent, child_name)

    def _check_loop(self, new_parent: Tree | None) -> None:
        """Checks that assignment of this new parent will not create a cycle."""
        if new_parent is not None:
            if new_parent is self:
                raise InvalidTreeError(
                    f"Cannot set parent, as node {self} cannot be a parent of itself."
                )

            if self._is_descendant_of(new_parent):
                raise InvalidTreeError(
                    "Cannot set parent, as intended parent is already a descendant of this node."
                )

    def _is_descendant_of(self, node: Tree) -> bool:
        return any(n is self for n in node.parents)

    def _detach(self, parent: Tree | None) -> None:
        if parent is not None:
            self._pre_detach(parent)
            parents_children = parent.children
            parent._children = {
                name: child
                for name, child in parents_children.items()
                if child is not self
            }
            self._parent = None
            self._post_detach(parent)

    def _attach(self, parent: Tree | None, child_name: str | None = None) -> None:
        if parent is not None:
            if child_name is None:
                raise ValueError(
                    "To directly set parent, child needs a name, but child is unnamed"
                )

            self._pre_attach(parent, child_name)
            parentchildren = parent._children
            assert not any(
                child is self for child in parentchildren
            ), "Tree is corrupt."
            parentchildren[child_name] = self
            self._parent = parent
            self._post_attach(parent, child_name)
        else:
            self._parent = None

    def orphan(self) -> None:
        """Detach this node from its parent."""
        self._set_parent(new_parent=None)

    @property
    def children(self: Tree) -> Mapping[str, Tree]:
        """Child nodes of this node, stored under a mapping via their names."""
        return Frozen(self._children)

    @children.setter
    def children(self: Tree, children: Mapping[str, Tree]) -> None:
        self._check_children(children)
        children = {**children}

        old_children = self.children
        del self.children
        try:
            self._pre_attach_children(children)
            for name, child in children.items():
                child._set_parent(new_parent=self, child_name=name)
            self._post_attach_children(children)
            assert len(self.children) == len(children)
        except Exception:
            # if something goes wrong then revert to previous children
            self.children = old_children
            raise

    @children.deleter
    def children(self) -> None:
        # TODO this just detaches all the children, it doesn't actually delete them...
        children = self.children
        self._pre_detach_children(children)
        for child in self.children.values():
            child.orphan()
        assert len(self.children) == 0
        self._post_detach_children(children)

    @staticmethod
    def _check_children(children: Mapping[str, Tree]) -> None:
        """Check children for correct types and for any duplicates."""
        if not is_dict_like(children):
            raise TypeError(
                "children must be a dict-like mapping from names to node objects"
            )

        seen = set()
        for name, child in children.items():
            if not isinstance(child, TreeNode):
                raise TypeError(
                    f"Cannot add object {name}. It is of type {type(child)}, "
                    "but can only add children of type DataTree"
                )

            childid = id(child)
            if childid not in seen:
                seen.add(childid)
            else:
                raise InvalidTreeError(
                    f"Cannot add same node {name} multiple times as different children."
                )

    def __repr__(self) -> str:
        return f"TreeNode(children={dict(self._children)})"

    def _pre_detach_children(self: Tree, children: Mapping[str, Tree]) -> None:
        """Method call before detaching `children`."""
        pass

    def _post_detach_children(self: Tree, children: Mapping[str, Tree]) -> None:
        """Method call after detaching `children`."""
        pass

    def _pre_attach_children(self: Tree, children: Mapping[str, Tree]) -> None:
        """Method call before attaching `children`."""
        pass

    def _post_attach_children(self: Tree, children: Mapping[str, Tree]) -> None:
        """Method call after attaching `children`."""
        pass

    def copy(
        self: Tree,
        deep: bool = False,
    ) -> Tree:
        """
        Returns a copy of this subtree.

        Copies this node and all child nodes.

        If `deep=True`, a deep copy is made of each of the component variables.
        Otherwise, a shallow copy of each of the component variable is made, so
        that the underlying memory region of the new datatree is the same as in
        the original datatree.

        Parameters
        ----------
        deep : bool, default: False
            Whether each component variable is loaded into memory and copied onto
            the new object. Default is False.

        Returns
        -------
        object : DataTree
            New object with dimensions, attributes, coordinates, name, encoding,
            and data of this node and all child nodes copied from original.

        See Also
        --------
        xarray.Dataset.copy
        pandas.DataFrame.copy
        """
        return self._copy_subtree(deep=deep)

    def _copy_subtree(
        self: Tree,
        deep: bool = False,
        memo: dict[int, Any] | None = None,
    ) -> Tree:
        """Copy entire subtree recursively."""

        new_tree = self._copy_node(deep=deep)
        for name, child in self.children.items():
            # TODO use `.children[name] = ...` once #9477 is implemented
            new_tree._set(name, child._copy_subtree(deep=deep))

        return new_tree

    def _copy_node(
        self: Tree,
        deep: bool = False,
    ) -> Tree:
        """Copy just one node of a tree"""
        new_empty_node = type(self)()
        return new_empty_node

    def __copy__(self: Tree) -> Tree:
        return self._copy_subtree(deep=False)

    def __deepcopy__(self: Tree, memo: dict[int, Any] | None = None) -> Tree:
        return self._copy_subtree(deep=True, memo=memo)

    def _iter_parents(self: Tree) -> Iterator[Tree]:
        """Iterate up the tree, starting from the current node's parent."""
        node: Tree | None = self.parent
        while node is not None:
            yield node
            node = node.parent

    def iter_lineage(self: Tree) -> tuple[Tree, ...]:
        """Iterate up the tree, starting from the current node."""
        from warnings import warn

        warn(
            "`iter_lineage` has been deprecated, and in the future will raise an error."
            "Please use `parents` from now on.",
            DeprecationWarning,
        )
        return tuple((self, *self.parents))

    @property
    def lineage(self: Tree) -> tuple[Tree, ...]:
        """All parent nodes and their parent nodes, starting with the closest."""
        from warnings import warn

        warn(
            "`lineage` has been deprecated, and in the future will raise an error."
            "Please use `parents` from now on.",
            DeprecationWarning,
        )
        return self.iter_lineage()

    @property
    def parents(self: Tree) -> tuple[Tree, ...]:
        """All parent nodes and their parent nodes, starting with the closest."""
        return tuple(self._iter_parents())

    @property
    def ancestors(self: Tree) -> tuple[Tree, ...]:
        """All parent nodes and their parent nodes, starting with the most distant."""

        from warnings import warn

        warn(
            "`ancestors` has been deprecated, and in the future will raise an error."
            "Please use `parents`. Example: `tuple(reversed(node.parents))`",
            DeprecationWarning,
        )
        return tuple((*reversed(self.parents), self))

    @property
    def root(self: Tree) -> Tree:
        """Root node of the tree"""
        node = self
        while node.parent is not None:
            node = node.parent
        return node

    @property
    def is_root(self) -> bool:
        """Whether this node is the tree root."""
        return self.parent is None

    @property
    def is_leaf(self) -> bool:
        """
        Whether this node is a leaf node.

        Leaf nodes are defined as nodes which have no children.
        """
        return self.children == {}

    @property
    def leaves(self: Tree) -> tuple[Tree, ...]:
        """
        All leaf nodes.

        Leaf nodes are defined as nodes which have no children.
        """
        return tuple([node for node in self.subtree if node.is_leaf])

    @property
    def siblings(self: Tree) -> dict[str, Tree]:
        """
        Nodes with the same parent as this node.
        """
        if self.parent:
            return {
                name: child
                for name, child in self.parent.children.items()
                if child is not self
            }
        else:
            return {}

    @property
    def subtree(self: Tree) -> Iterator[Tree]:
        """
        An iterator over all nodes in this tree, including both self and all descendants.

        Iterates depth-first.

        See Also
        --------
        DataTree.descendants
        """
        from xarray.core.iterators import LevelOrderIter

        return LevelOrderIter(self)

    @property
    def descendants(self: Tree) -> tuple[Tree, ...]:
        """
        Child nodes and all their child nodes.

        Returned in depth-first order.

        See Also
        --------
        DataTree.subtree
        """
        all_nodes = tuple(self.subtree)
        this_node, *descendants = all_nodes
        return tuple(descendants)

    @property
    def level(self: Tree) -> int:
        """
        Level of this node.

        Level means number of parent nodes above this node before reaching the root.
        The root node is at level 0.

        Returns
        -------
        level : int

        See Also
        --------
        depth
        width
        """
        return len(self.parents)

    @property
    def depth(self: Tree) -> int:
        """
        Maximum level of this tree.

        Measured from the root, which has a depth of 0.

        Returns
        -------
        depth : int

        See Also
        --------
        level
        width
        """
        return max(node.level for node in self.root.subtree)

    @property
    def width(self: Tree) -> int:
        """
        Number of nodes at this level in the tree.

        Includes number of immediate siblings, but also "cousins" in other branches and so-on.

        Returns
        -------
        depth : int

        See Also
        --------
        level
        depth
        """
        return len([node for node in self.root.subtree if node.level == self.level])

    def _pre_detach(self: Tree, parent: Tree) -> None:
        """Method call before detaching from `parent`."""
        pass

    def _post_detach(self: Tree, parent: Tree) -> None:
        """Method call after detaching from `parent`."""
        pass

    def _pre_attach(self: Tree, parent: Tree, name: str) -> None:
        """Method call before attaching to `parent`."""
        pass

    def _post_attach(self: Tree, parent: Tree, name: str) -> None:
        """Method call after attaching to `parent`."""
        pass

    def get(self: Tree, key: str, default: Tree | None = None) -> Tree | None:
        """
        Return the child node with the specified key.

        Only looks for the node within the immediate children of this node,
        not in other nodes of the tree.
        """
        if key in self.children:
            return self.children[key]
        else:
            return default

    # TODO `._walk` method to be called by both `_get_item` and `_set_item`

    def _get_item(self: Tree, path: str | NodePath) -> Tree | T_DataArray:
        """
        Returns the object lying at the given path.

        Raises a KeyError if there is no object at the given path.
        """
        if isinstance(path, str):
            path = NodePath(path)

        if path.root:
            current_node = self.root
            root, *parts = list(path.parts)
        else:
            current_node = self
            parts = list(path.parts)

        for part in parts:
            if part == "..":
                if current_node.parent is None:
                    raise KeyError(f"Could not find node at {path}")
                else:
                    current_node = current_node.parent
            elif part in ("", "."):
                pass
            else:
                if current_node.get(part) is None:
                    raise KeyError(f"Could not find node at {path}")
                else:
                    current_node = current_node.get(part)
        return current_node

    def _set(self: Tree, key: str, val: Tree) -> None:
        """
        Set the child node with the specified key to value.

        Counterpart to the public .get method, and also only works on the immediate node, not other nodes in the tree.
        """
        new_children = {**self.children, key: val}
        self.children = new_children

    def _set_item(
        self: Tree,
        path: str | NodePath,
        item: Tree | T_DataArray,
        new_nodes_along_path: bool = False,
        allow_overwrite: bool = True,
    ) -> None:
        """
        Set a new item in the tree, overwriting anything already present at that path.

        The given value either forms a new node of the tree or overwrites an
        existing item at that location.

        Parameters
        ----------
        path
        item
        new_nodes_along_path : bool
            If true, then if necessary new nodes will be created along the
            given path, until the tree can reach the specified location.
        allow_overwrite : bool
            Whether or not to overwrite any existing node at the location given
            by path.

        Raises
        ------
        KeyError
            If node cannot be reached, and new_nodes_along_path=False.
            Or if a node already exists at the specified path, and allow_overwrite=False.
        """
        if isinstance(path, str):
            path = NodePath(path)

        if not path.name:
            raise ValueError("Can't set an item under a path which has no name")

        if path.root:
            # absolute path
            current_node = self.root
            root, *parts, name = path.parts
        else:
            # relative path
            current_node = self
            *parts, name = path.parts

        if parts:
            # Walk to location of new node, creating intermediate node objects as we go if necessary
            for part in parts:
                if part == "..":
                    if current_node.parent is None:
                        # We can't create a parent if `new_nodes_along_path=True` as we wouldn't know what to name it
                        raise KeyError(f"Could not reach node at path {path}")
                    else:
                        current_node = current_node.parent
                elif part in ("", "."):
                    pass
                else:
                    if part in current_node.children:
                        current_node = current_node.children[part]
                    elif new_nodes_along_path:
                        # Want child classes (i.e. DataTree) to populate tree with their own types
                        new_node = type(self)()
                        current_node._set(part, new_node)
                        current_node = current_node.children[part]
                    else:
                        raise KeyError(f"Could not reach node at path {path}")

        if name in current_node.children:
            # Deal with anything already existing at this location
            if allow_overwrite:
                current_node._set(name, item)
            else:
                raise KeyError(f"Already a node object at path {path}")
        else:
            current_node._set(name, item)

    def __delitem__(self: Tree, key: str) -> None:
        """Remove a child node from this tree object."""
        if key in self.children:
            child = self._children[key]
            del self._children[key]
            child.orphan()
        else:
            raise KeyError(key)

    def same_tree(self, other: Tree) -> bool:
        """True if other node is in the same tree as this node."""
        return self.root is other.root


AnyNamedNode = TypeVar("AnyNamedNode", bound="NamedNode")


class NamedNode(TreeNode, Generic[Tree]):
    """
    A TreeNode which knows its own name.

    Implements path-like relationships to other nodes in its tree.
    """

    _name: str | None
    _parent: Tree | None
    _children: dict[str, Tree]

    def __init__(self, name=None, children=None):
        super().__init__(children=children)
        self._name = None
        self.name = name

    @property
    def name(self) -> str | None:
        """The name of this node."""
        return self._name

    @name.setter
    def name(self, name: str | None) -> None:
        if name is not None:
            if not isinstance(name, str):
                raise TypeError("node name must be a string or None")
            if "/" in name:
                raise ValueError("node names cannot contain forward slashes")
        self._name = name

    def __repr__(self, level=0):
        repr_value = "\t" * level + self.__str__() + "\n"
        for child in self.children:
            repr_value += self.get(child).__repr__(level + 1)
        return repr_value

    def __str__(self) -> str:
        return f"NamedNode('{self.name}')" if self.name else "NamedNode()"

    def _post_attach(self: AnyNamedNode, parent: AnyNamedNode, name: str) -> None:
        """Ensures child has name attribute corresponding to key under which it has been stored."""
        self.name = name

    def _copy_node(
        self: AnyNamedNode,
        deep: bool = False,
    ) -> AnyNamedNode:
        """Copy just one node of a tree"""
        new_node = super()._copy_node()
        new_node._name = self.name
        return new_node

    @property
    def path(self) -> str:
        """Return the file-like path from the root to this node."""
        if self.is_root:
            return "/"
        else:
            root, *ancestors = tuple(reversed(self.parents))
            # don't include name of root because (a) root might not have a name & (b) we want path relative to root.
            names = [*(node.name for node in ancestors), self.name]
            return "/" + "/".join(names)

    def relative_to(self: NamedNode, other: NamedNode) -> str:
        """
        Compute the relative path from this node to node `other`.

        If other is not in this tree, or it's otherwise impossible, raise a ValueError.
        """
        if not self.same_tree(other):
            raise NotFoundInTreeError(
                "Cannot find relative path because nodes do not lie within the same tree"
            )

        this_path = NodePath(self.path)
        if other.path in list(parent.path for parent in (self, *self.parents)):
            return str(this_path.relative_to(other.path))
        else:
            common_ancestor = self.find_common_ancestor(other)
            path_to_common_ancestor = other._path_to_ancestor(common_ancestor)
            return str(
                path_to_common_ancestor / this_path.relative_to(common_ancestor.path)
            )

    def find_common_ancestor(self, other: NamedNode) -> NamedNode:
        """
        Find the first common ancestor of two nodes in the same tree.

        Raise ValueError if they are not in the same tree.
        """
        if self is other:
            return self

        other_paths = [op.path for op in other.parents]
        for parent in (self, *self.parents):
            if parent.path in other_paths:
                return parent

        raise NotFoundInTreeError(
            "Cannot find common ancestor because nodes do not lie within the same tree"
        )

    def _path_to_ancestor(self, ancestor: NamedNode) -> NodePath:
        """Return the relative path from this node to the given ancestor node"""

        if not self.same_tree(ancestor):
            raise NotFoundInTreeError(
                "Cannot find relative path to ancestor because nodes do not lie within the same tree"
            )
        if ancestor.path not in list(a.path for a in (self, *self.parents)):
            raise NotFoundInTreeError(
                "Cannot find relative path to ancestor because given node is not an ancestor of this node"
            )

        parents_paths = list(parent.path for parent in (self, *self.parents))
        generation_gap = list(parents_paths).index(ancestor.path)
        path_upwards = "../" * generation_gap if generation_gap > 0 else "."
        return NodePath(path_upwards)