matcher.py 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. class MQTTMatcher(object):
  2. """Intended to manage topic filters including wildcards.
  3. Internally, MQTTMatcher use a prefix tree (trie) to store
  4. values associated with filters, and has an iter_match()
  5. method to iterate efficiently over all filters that match
  6. some topic name."""
  7. class Node(object):
  8. __slots__ = '_children', '_content'
  9. def __init__(self):
  10. self._children = {}
  11. self._content = None
  12. def __init__(self):
  13. self._root = self.Node()
  14. def __setitem__(self, key, value):
  15. """Add a topic filter :key to the prefix tree
  16. and associate it to :value"""
  17. node = self._root
  18. for sym in key.split('/'):
  19. node = node._children.setdefault(sym, self.Node())
  20. node._content = value
  21. def __getitem__(self, key):
  22. """Retrieve the value associated with some topic filter :key"""
  23. try:
  24. node = self._root
  25. for sym in key.split('/'):
  26. node = node._children[sym]
  27. if node._content is None:
  28. raise KeyError(key)
  29. return node._content
  30. except KeyError:
  31. raise KeyError(key)
  32. def __delitem__(self, key):
  33. """Delete the value associated with some topic filter :key"""
  34. lst = []
  35. try:
  36. parent, node = None, self._root
  37. for k in key.split('/'):
  38. parent, node = node, node._children[k]
  39. lst.append((parent, k, node))
  40. # TODO
  41. node._content = None
  42. except KeyError:
  43. raise KeyError(key)
  44. else: # cleanup
  45. for parent, k, node in reversed(lst):
  46. if node._children or node._content is not None:
  47. break
  48. del parent._children[k]
  49. def iter_match(self, topic):
  50. """Return an iterator on all values associated with filters
  51. that match the :topic"""
  52. lst = topic.split('/')
  53. normal = not topic.startswith('$')
  54. def rec(node, i=0):
  55. if i == len(lst):
  56. if node._content is not None:
  57. yield node._content
  58. else:
  59. part = lst[i]
  60. if part in node._children:
  61. for content in rec(node._children[part], i + 1):
  62. yield content
  63. if '+' in node._children and (normal or i > 0):
  64. for content in rec(node._children['+'], i + 1):
  65. yield content
  66. if '#' in node._children and (normal or i > 0):
  67. content = node._children['#']._content
  68. if content is not None:
  69. yield content
  70. return rec(self._root)