tree.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. """
  2. A class for storing a tree graph. Primarily used for filter constructs in the
  3. ORM.
  4. """
  5. import copy
  6. from django.utils.hashable import make_hashable
  7. class Node:
  8. """
  9. A single internal node in the tree graph. A Node should be viewed as a
  10. connection (the root) with the children being either leaf nodes or other
  11. Node instances.
  12. """
  13. # Standard connector type. Clients usually won't use this at all and
  14. # subclasses will usually override the value.
  15. default = "DEFAULT"
  16. def __init__(self, children=None, connector=None, negated=False):
  17. """Construct a new Node. If no connector is given, use the default."""
  18. self.children = children[:] if children else []
  19. self.connector = connector or self.default
  20. self.negated = negated
  21. @classmethod
  22. def create(cls, children=None, connector=None, negated=False):
  23. """
  24. Create a new instance using Node() instead of __init__() as some
  25. subclasses, e.g. django.db.models.query_utils.Q, may implement a custom
  26. __init__() with a signature that conflicts with the one defined in
  27. Node.__init__().
  28. """
  29. obj = Node(children, connector or cls.default, negated)
  30. obj.__class__ = cls
  31. return obj
  32. def __str__(self):
  33. template = "(NOT (%s: %s))" if self.negated else "(%s: %s)"
  34. return template % (self.connector, ", ".join(str(c) for c in self.children))
  35. def __repr__(self):
  36. return "<%s: %s>" % (self.__class__.__name__, self)
  37. def __copy__(self):
  38. obj = self.create(connector=self.connector, negated=self.negated)
  39. obj.children = self.children # Don't [:] as .__init__() via .create() does.
  40. return obj
  41. copy = __copy__
  42. def __deepcopy__(self, memodict):
  43. obj = self.create(connector=self.connector, negated=self.negated)
  44. obj.children = copy.deepcopy(self.children, memodict)
  45. return obj
  46. def __len__(self):
  47. """Return the number of children this node has."""
  48. return len(self.children)
  49. def __bool__(self):
  50. """Return whether or not this node has children."""
  51. return bool(self.children)
  52. def __contains__(self, other):
  53. """Return True if 'other' is a direct child of this instance."""
  54. return other in self.children
  55. def __eq__(self, other):
  56. return (
  57. self.__class__ == other.__class__
  58. and self.connector == other.connector
  59. and self.negated == other.negated
  60. and self.children == other.children
  61. )
  62. def __hash__(self):
  63. return hash(
  64. (
  65. self.__class__,
  66. self.connector,
  67. self.negated,
  68. *make_hashable(self.children),
  69. )
  70. )
  71. def add(self, data, conn_type):
  72. """
  73. Combine this tree and the data represented by data using the
  74. connector conn_type. The combine is done by squashing the node other
  75. away if possible.
  76. This tree (self) will never be pushed to a child node of the
  77. combined tree, nor will the connector or negated properties change.
  78. Return a node which can be used in place of data regardless if the
  79. node other got squashed or not.
  80. """
  81. if self.connector != conn_type:
  82. obj = self.copy()
  83. self.connector = conn_type
  84. self.children = [obj, data]
  85. return data
  86. elif (
  87. isinstance(data, Node)
  88. and not data.negated
  89. and (data.connector == conn_type or len(data) == 1)
  90. ):
  91. # We can squash the other node's children directly into this node.
  92. # We are just doing (AB)(CD) == (ABCD) here, with the addition that
  93. # if the length of the other node is 1 the connector doesn't
  94. # matter. However, for the len(self) == 1 case we don't want to do
  95. # the squashing, as it would alter self.connector.
  96. self.children.extend(data.children)
  97. return self
  98. else:
  99. # We could use perhaps additional logic here to see if some
  100. # children could be used for pushdown here.
  101. self.children.append(data)
  102. return data
  103. def negate(self):
  104. """Negate the sense of the root connector."""
  105. self.negated = not self.negated