django/utils/tree.py

"""
A class for storing a tree graph. Primarily used for filter constructs in the
ORM.
"""

import copy

from django.utils.hashable import make_hashable


class Node:
    """
    A single internal node in the tree graph. A Node should be viewed as a
    connection (the root) with the children being either leaf nodes or other
    Node instances.
    """

    # Standard connector type. Clients usually won't use this at all and
    # subclasses will usually override the value.
    default = "DEFAULT"

    def __init__(self, children=None, connector=None, negated=False):
        """Construct a new Node. If no connector is given, use the default."""
        self.children = children[:] if children else []
        self.connector = connector or self.default
        self.negated = negated

    @classmethod
    def create(cls, children=None, connector=None, negated=False):
        """
        Create a new instance using Node() instead of __init__() as some
        subclasses, e.g. django.db.models.query_utils.Q, may implement a custom
        __init__() with a signature that conflicts with the one defined in
        Node.__init__().
        """
        obj = Node(children, connector or cls.default, negated)
        obj.__class__ = cls
        return obj

    def __str__(self):
        template = "(NOT (%s: %s))" if self.negated else "(%s: %s)"
        return template % (self.connector, ", ".join(str(c) for c in self.children))

    def __repr__(self):
        return "<%s: %s>" % (self.__class__.__name__, self)

    def __copy__(self):
        obj = self.create(connector=self.connector, negated=self.negated)
        obj.children = self.children  # Don't [:] as .__init__() via .create() does.
        return obj

    copy = __copy__

    def __deepcopy__(self, memodict):
        obj = self.create(connector=self.connector, negated=self.negated)
        obj.children = copy.deepcopy(self.children, memodict)
        return obj

    def __len__(self):
        """Return the number of children this node has."""
        return len(self.children)

    def __bool__(self):
        """Return whether or not this node has children."""
        return bool(self.children)

    def __contains__(self, other):
        """Return True if 'other' is a direct child of this instance."""
        return other in self.children

    def __eq__(self, other):
        return (
            self.__class__ == other.__class__
            and self.connector == other.connector
            and self.negated == other.negated
            and self.children == other.children
        )

    def __hash__(self):
        return hash(
            (
                self.__class__,
                self.connector,
                self.negated,
                *make_hashable(self.children),
            )
        )

    def add(self, data, conn_type):
        """
        Combine this tree and the data represented by data using the
        connector conn_type. The combine is done by squashing the node other
        away if possible.

        This tree (self) will never be pushed to a child node of the
        combined tree, nor will the connector or negated properties change.

        Return a node which can be used in place of data regardless if the
        node other got squashed or not.
        """
        if self.connector != conn_type:
            obj = self.copy()
            self.connector = conn_type
            self.children = [obj, data]
            return data
        elif (
            isinstance(data, Node)
            and not data.negated
            and (data.connector == conn_type or len(data) == 1)
        ):
            # We can squash the other node's children directly into this node.
            # We are just doing (AB)(CD) == (ABCD) here, with the addition that
            # if the length of the other node is 1 the connector doesn't
            # matter. However, for the len(self) == 1 case we don't want to do
            # the squashing, as it would alter self.connector.
            self.children.extend(data.children)
            return self
        else:
            # We could use perhaps additional logic here to see if some
            # children could be used for pushdown here.
            self.children.append(data)
            return data

    def negate(self):
        """Negate the sense of the root connector."""
        self.negated = not self.negated