django/utils/tree.py

   1 """
   2 A class for storing a tree graph. Primarily used for filter constructs in the
   3 ORM.
   4 """
   5
   6 from copy import deepcopy
   7
   8 class Node(object):
   9     """
  10     A single internal node in the tree graph. A Node should be viewed as a
  11     connection (the root) with the children being either leaf nodes or other
  12     Node instances.
  13     """
  14     # Standard connector type. Clients usually won't use this at all and
  15     # subclasses will usually override the value.
  16     default = 'DEFAULT'
  17
  18     def __init__(self, children=None, connector=None, negated=False):
  19         """
  20         Constructs a new Node. If no connector is given, the default will be
  21         used.
  22
  23         Warning: You probably don't want to pass in the 'negated' parameter. It
  24         is NOT the same as constructing a node and calling negate() on the
  25         result.
  26         """
  27         self.children = children and children[:] or []
  28         self.connector = connector or self.default
  29         self.subtree_parents = []
  30         self.negated = negated
  31
  32     # We need this because of django.db.models.query_utils.Q. Q. __init__() is
  33     # problematic, but it is a natural Node subclass in all other respects.
  34     def _new_instance(cls, children=None, connector=None, negated=False):
  35         """
  36         This is called to create a new instance of this class when we need new
  37         Nodes (or subclasses) in the internal code in this class. Normally, it
  38         just shadows __init__(). However, subclasses with an __init__ signature
  39         that is not an extension of Node.__init__ might need to implement this
  40         method to allow a Node to create a new instance of them (if they have
  41         any extra setting up to do).
  42         """
  43         obj = Node(children, connector, negated)
  44         obj.__class__ = cls
  45         return obj
  46     _new_instance = classmethod(_new_instance)
  47
  48     def __str__(self):
  49         if self.negated:
  50             return '(NOT (%s: %s))' % (self.connector, ', '.join([str(c) for c
  51                     in self.children]))
  52         return '(%s: %s)' % (self.connector, ', '.join([str(c) for c in
  53                 self.children]))
  54
  55     def __deepcopy__(self, memodict):
  56         """
  57         Utility method used by copy.deepcopy().
  58         """
  59         obj = Node(connector=self.connector, negated=self.negated)
  60         obj.__class__ = self.__class__
  61         obj.children = deepcopy(self.children, memodict)
  62         obj.subtree_parents = deepcopy(self.subtree_parents, memodict)
  63         return obj
  64
  65     def __len__(self):
  66         """
  67         The size of a node if the number of children it has.
  68         """
  69         return len(self.children)
  70
  71     def __nonzero__(self):
  72         """
  73         For truth value testing.
  74         """
  75         return bool(self.children)
  76
  77     def __contains__(self, other):
  78         """
  79         Returns True is 'other' is a direct child of this instance.
  80         """
  81         return other in self.children
  82
  83     def add(self, node, conn_type):
  84         """
  85         Adds a new node to the tree. If the conn_type is the same as the root's
  86         current connector type, the node is added to the first level.
  87         Otherwise, the whole tree is pushed down one level and a new root
  88         connector is created, connecting the existing tree and the new node.
  89         """
  90         if node in self.children and conn_type == self.connector:
  91             return
  92         if len(self.children) < 2:
  93             self.connector = conn_type
  94         if self.connector == conn_type:
  95             if isinstance(node, Node) and (node.connector == conn_type or
  96                     len(node) == 1):
  97                 self.children.extend(node.children)
  98             else:
  99                 self.children.append(node)
 100         else:
 101             obj = self._new_instance(self.children, self.connector,
 102                     self.negated)
 103             self.connector = conn_type
 104             self.children = [obj, node]
 105
 106     def negate(self):
 107         """
 108         Negate the sense of the root connector. This reorganises the children
 109         so that the current node has a single child: a negated node containing
 110         all the previous children. This slightly odd construction makes adding
 111         new children behave more intuitively.
 112
 113         Interpreting the meaning of this negate is up to client code. This
 114         method is useful for implementing "not" arrangements.
 115         """
 116         self.children = [self._new_instance(self.children, self.connector,
 117                 not self.negated)]
 118         self.connector = self.default
 119
 120     def start_subtree(self, conn_type):
 121         """
 122         Sets up internal state so that new nodes are added to a subtree of the
 123         current node. The conn_type specifies how the sub-tree is joined to the
 124         existing children.
 125         """
 126         if len(self.children) == 1:
 127             self.connector = conn_type
 128         elif self.connector != conn_type:
 129             self.children = [self._new_instance(self.children, self.connector,
 130                     self.negated)]
 131             self.connector = conn_type
 132             self.negated = False
 133
 134         self.subtree_parents.append(self.__class__(self.children,
 135                 self.connector, self.negated))
 136         self.connector = self.default
 137         self.negated = False
 138         self.children = []
 139
 140     def end_subtree(self):
 141         """
 142         Closes off the most recently unmatched start_subtree() call.
 143
 144         This puts the current state into a node of the parent tree and returns
 145         the current instances state to be the parent.
 146         """
 147         obj = self.subtree_parents.pop()
 148         node = self.__class__(self.children, self.connector)
 149         self.connector = obj.connector
 150         self.negated = obj.negated
 151         self.children = obj.children
 152         self.children.append(node)
 153