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     def __str__(self):
  33         if self.negated:
  34             return '(NOT (%s: %s))' % (self.connector, ', '.join([str(c) for c
  35                     in self.children]))
  36         return '(%s: %s)' % (self.connector, ', '.join([str(c) for c in
  37                 self.children]))
  38
  39     def __deepcopy__(self, memodict):
  40         """
  41         Utility method used by copy.deepcopy().
  42         """
  43         obj = Node(connector=self.connector, negated=self.negated)
  44         obj.__class__ = self.__class__
  45         obj.children = deepcopy(self.children, memodict)
  46         obj.subtree_parents = deepcopy(self.subtree_parents, memodict)
  47         return obj
  48
  49     def __len__(self):
  50         """
  51         The size of a node if the number of children it has.
  52         """
  53         return len(self.children)
  54
  55     def __nonzero__(self):
  56         """
  57         For truth value testing.
  58         """
  59         return bool(self.children)
  60
  61     def __contains__(self, other):
  62         """
  63         Returns True is 'other' is a direct child of this instance.
  64         """
  65         return other in self.children
  66
  67     def add(self, node, conn_type):
  68         """
  69         Adds a new node to the tree. If the conn_type is the same as the root's
  70         current connector type, the node is added to the first level.
  71         Otherwise, the whole tree is pushed down one level and a new root
  72         connector is created, connecting the existing tree and the new node.
  73         """
  74         if node in self.children:
  75             return
  76         if len(self.children) < 2:
  77             self.connector = conn_type
  78         if self.connector == conn_type:
  79             if isinstance(node, Node) and (node.connector == conn_type or
  80                     len(node) == 1):
  81                 self.children.extend(node.children)
  82             else:
  83                 self.children.append(node)
  84         else:
  85             obj = Node(self.children, self.connector, self.negated)
  86             self.connector = conn_type
  87             self.children = [obj, node]
  88
  89     def negate(self):
  90         """
  91         Negate the sense of the root connector. This reorganises the children
  92         so that the current node has a single child: a negated node containing
  93         all the previous children. This slightly odd construction makes adding
  94         new children behave more intuitively.
  95
  96         Interpreting the meaning of this negate is up to client code. This
  97         method is useful for implementing "not" arrangements.
  98         """
  99         self.children = [Node(self.children, self.connector, not self.negated)]
 100         self.connector = self.default
 101
 102     def start_subtree(self, conn_type):
 103         """
 104         Sets up internal state so that new nodes are added to a subtree of the
 105         current node. The conn_type specifies how the sub-tree is joined to the
 106         existing children.
 107         """
 108         if len(self.children) == 1:
 109             self.connector = conn_type
 110         elif self.connector != conn_type:
 111             self.children = [Node(self.children, self.connector, self.negated)]
 112             self.connector = conn_type
 113             self.negated = False
 114
 115         self.subtree_parents.append(Node(self.children, self.connector,
 116                 self.negated))
 117         self.connector = self.default
 118         self.negated = False
 119         self.children = []
 120
 121     def end_subtree(self):
 122         """
 123         Closes off the most recently unmatched start_subtree() call.
 124
 125         This puts the current state into a node of the parent tree and returns
 126         the current instances state to be the parent.
 127         """
 128         obj = self.subtree_parents.pop()
 129         node = Node(self.children, self.connector)
 130         self.connector = obj.connector
 131         self.negated = obj.negated
 132         self.children = obj.children
 133         self.children.append(node)
 134