tree.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. """
  2. A class for storing a tree graph. Primarily used for filter constructs in the
  3. ORM.
  4. """
  5. import copy
  6. class Node(object):
  7. """
  8. A single internal node in the tree graph. A Node should be viewed as a
  9. connection (the root) with the children being either leaf nodes or other
  10. Node instances.
  11. """
  12. # Standard connector type. Clients usually won't use this at all and
  13. # subclasses will usually override the value.
  14. default = 'DEFAULT'
  15. def __init__(self, children=None, connector=None, negated=False):
  16. """
  17. Constructs a new Node. If no connector is given, the default will be
  18. used.
  19. """
  20. self.children = children[:] if children else []
  21. self.connector = connector or self.default
  22. self.negated = negated
  23. # We need this because of django.db.models.query_utils.Q. Q. __init__() is
  24. # problematic, but it is a natural Node subclass in all other respects.
  25. @classmethod
  26. def _new_instance(cls, children=None, connector=None, negated=False):
  27. """
  28. This is called to create a new instance of this class when we need new
  29. Nodes (or subclasses) in the internal code in this class. Normally, it
  30. just shadows __init__(). However, subclasses with an __init__ signature
  31. that is not an extension of Node.__init__ might need to implement this
  32. method to allow a Node to create a new instance of them (if they have
  33. any extra setting up to do).
  34. """
  35. obj = Node(children, connector, negated)
  36. obj.__class__ = cls
  37. return obj
  38. def __str__(self):
  39. if self.negated:
  40. return '(NOT (%s: %s))' % (self.connector, ', '.join([str(c) for c
  41. in self.children]))
  42. return '(%s: %s)' % (self.connector, ', '.join([str(c) for c in
  43. self.children]))
  44. def __deepcopy__(self, memodict):
  45. """
  46. Utility method used by copy.deepcopy().
  47. """
  48. obj = Node(connector=self.connector, negated=self.negated)
  49. obj.__class__ = self.__class__
  50. obj.children = copy.deepcopy(self.children, memodict)
  51. return obj
  52. def __len__(self):
  53. """
  54. The size of a node if the number of children it has.
  55. """
  56. return len(self.children)
  57. def __bool__(self):
  58. """
  59. For truth value testing.
  60. """
  61. return bool(self.children)
  62. def __nonzero__(self): # Python 2 compatibility
  63. return type(self).__bool__(self)
  64. def __contains__(self, other):
  65. """
  66. Returns True is 'other' is a direct child of this instance.
  67. """
  68. return other in self.children
  69. def _prepare_data(self, data):
  70. """
  71. A subclass hook for doing subclass specific transformations of the
  72. given data on combine() or add().
  73. """
  74. return data
  75. def add(self, data, conn_type, squash=True):
  76. """
  77. Combines this tree and the data represented by data using the
  78. connector conn_type. The combine is done by squashing the node other
  79. away if possible.
  80. This tree (self) will never be pushed to a child node of the
  81. combined tree, nor will the connector or negated properties change.
  82. The function returns a node which can be used in place of data
  83. regardless if the node other got squashed or not.
  84. If `squash` is False the data is prepared and added as a child to
  85. this tree without further logic.
  86. """
  87. if data in self.children:
  88. return data
  89. data = self._prepare_data(data)
  90. if not squash:
  91. self.children.append(data)
  92. return data
  93. if self.connector == conn_type:
  94. # We can reuse self.children to append or squash the node other.
  95. if (isinstance(data, Node) and not data.negated
  96. and (data.connector == conn_type or len(data) == 1)):
  97. # We can squash the other node's children directly into this
  98. # node. We are just doing (AB)(CD) == (ABCD) here, with the
  99. # addition that if the length of the other node is 1 the
  100. # connector doesn't matter. However, for the len(self) == 1
  101. # case we don't want to do the squashing, as it would alter
  102. # self.connector.
  103. self.children.extend(data.children)
  104. return self
  105. else:
  106. # We could use perhaps additional logic here to see if some
  107. # children could be used for pushdown here.
  108. self.children.append(data)
  109. return data
  110. else:
  111. obj = self._new_instance(self.children, self.connector,
  112. self.negated)
  113. self.connector = conn_type
  114. self.children = [obj, data]
  115. return data
  116. def negate(self):
  117. """
  118. Negate the sense of the root connector.
  119. """
  120. self.negated = not self.negated