add MQTTMatcher, expose MQTT's on_message, allow multiple callbacks

brentru · brentru · commit 9fad3122a366 · 2020-06-05T15:25:20.000-04:00
diff --git a/adafruit_minimqtt/__init__.py b/adafruit_minimqtt/__init__.py
diff --git a/adafruit_minimqtt/adafruit_minimqtt.py b/adafruit_minimqtt/adafruit_minimqtt.py
@@ -43,6 +43,7 @@
 import time
 from random import randint
 from micropython import const
+from .matcher import MQTTMatcher
 import adafruit_logging as logging
 
 __version__ = "0.0.0-auto.0"
@@ -173,8 +174,9 @@ def __init__(
         self._lw_retain = False
         # List of subscribed topics, used for tracking
         self._subscribed_topics = []
+        self._on_message_filtered = MQTTMatcher()
         # Server callbacks
-        self.on_message = None
+        self._on_message = None
         self.on_connect = None
         self.on_disconnect = None
         self.on_publish = None
@@ -218,6 +220,49 @@ def will_set(self, topic=None, payload=None, qos=0, retain=False):
         self._lw_msg = payload
         self._lw_retain = retain
 
+    def add_topic_callback(self, mqtt_topic, callback_method):
+        """Registers a callback_method for a specific MQTT topic.
+        :param str mqtt_topic: MQTT topic.
+        :param str callback_method: Name of callback method.
+
+        """
+        print("adding topic callback...")
+        if mqtt_topic is None or callback_method is None:
+            raise ValueError("MQTT topic and callback method must both be defined.")
+        self._on_message_filtered[mqtt_topic] = callback_method
+
+    def remove_topic_callback(self, mqtt_topic):
+        """Removes a registered callback method.
+        :param str mqtt_topic: MQTT topic.
+
+        """
+        if mqtt_topic is None:
+            raise ValueError("MQTT Topic must be defined.")
+        pass
+
+    @property
+    def on_message(self):
+        """Called when a new message has been received on a subscribed topic.
+
+        Expected method signature is:
+            on_message(client, topic, message)
+        """
+        return self._on_message
+
+    @on_message.setter
+    def on_message(self, method):
+        self._on_message = method
+
+    def _handle_on_message(self, client, topic, message):
+        matched = False
+        if topic is not None:
+            for callback in self._on_message_filtered.iter_match(topic):
+                callback(client, topic, message) # on_msg with callback
+                matched = True
+
+        if matched == False and self.on_message: # regular on_message
+            self.on_message(client, topic, message)
+
     # pylint: disable=too-many-branches, too-many-statements, too-many-locals
     def connect(self, clean_session=True):
         """Initiates connection with the MQTT Broker.
@@ -666,7 +711,7 @@ def _wait_for_msg(self, timeout=30):
             sz -= 0x02
         msg = self._sock.recv(sz)
         if self.on_message is not None:
-            self.on_message(self, topic, str(msg, "utf-8"))
+            self._handle_on_message(self, topic, str(msg, "utf-8"))
         if res[0] & 0x06 == 0x02:
             pkt = bytearray(b"\x40\x02\0\0")
             struct.pack_into("!H", pkt, 2, pid)
@@ -751,7 +796,7 @@ def mqtt_msg(self, msg_size):
         if msg_size < MQTT_MSG_MAX_SZ:
             self._msg_size_lim = msg_size
 
-    # Logging
+    ### Logging ###
     def attach_logger(self, logger_name="log"):
         """Initializes and attaches a logger to the MQTTClient.
         :param str logger_name: Name of the logger instance
diff --git a/adafruit_minimqtt/matcher.py b/adafruit_minimqtt/matcher.py
@@ -0,0 +1,80 @@
+class MQTTMatcher(object):
+    """Intended to manage topic filters including wildcards.
+
+    Internally, MQTTMatcher use a prefix tree (trie) to store 
+    values associated with filters, and has an iter_match() 
+    method to iterate efficiently over all filters that match 
+    some topic name."""
+
+    class Node(object):
+        __slots__ = "_children", "_content"
+
+        def __init__(self):
+            self._children = {}
+            self._content = None
+
+    def __init__(self):
+        self._root = self.Node()
+
+    def __setitem__(self, key, value):
+        """Add a topic filter :key to the prefix tree
+        and associate it to :value"""
+        node = self._root
+        for sym in key.split("/"):
+            node = node._children.setdefault(sym, self.Node())
+        node._content = value
+
+    def __getitem__(self, key):
+        """Retrieve the value associated with some topic filter :key"""
+        try:
+            node = self._root
+            for sym in key.split("/"):
+                node = node._children[sym]
+            if node._content is None:
+                raise KeyError(key)
+            return node._content
+        except KeyError:
+            raise KeyError(key)
+
+    def __delitem__(self, key):
+        """Delete the value associated with some topic filter :key"""
+        lst = []
+        try:
+            parent, node = None, self._root
+            for k in key.split("/"):
+                parent, node = node, node._children[k]
+                lst.append((parent, k, node))
+            # TODO
+            node._content = None
+        except KeyError:
+            raise KeyError(key)
+        else:  # cleanup
+            for parent, k, node in reversed(lst):
+                if node._children or node._content is not None:
+                    break
+                del parent._children[k]
+
+    def iter_match(self, topic):
+        """Return an iterator on all values associated with filters 
+        that match the :topic"""
+        lst = topic.split("/")
+        normal = not topic.startswith("$")
+
+        def rec(node, i=0):
+            if i == len(lst):
+                if node._content is not None:
+                    yield node._content
+            else:
+                part = lst[i]
+                if part in node._children:
+                    for content in rec(node._children[part], i + 1):
+                        yield content
+                if "+" in node._children and (normal or i > 0):
+                    for content in rec(node._children["+"], i + 1):
+                        yield content
+            if "#" in node._children and (normal or i > 0):
+                content = node._children["#"]._content
+                if content is not None:
+                    yield content
+
+        return rec(self._root)
diff --git a/examples/minimqtt_pub_sub_multiple_callbacks.py b/examples/minimqtt_pub_sub_multiple_callbacks.py
@@ -0,0 +1,135 @@
+import time
+import board
+import busio
+from digitalio import DigitalInOut
+import neopixel
+from adafruit_esp32spi import adafruit_esp32spi
+from adafruit_esp32spi import adafruit_esp32spi_wifimanager
+import adafruit_esp32spi.adafruit_esp32spi_socket as socket
+import adafruit_minimqtt.adafruit_minimqtt as MQTT
+
+### WiFi ###
+
+# Get wifi details and more from a secrets.py file
+try:
+    from secrets import secrets
+except ImportError:
+    print("WiFi secrets are kept in secrets.py, please add them there!")
+    raise
+
+esp32_cs = DigitalInOut(board.D13)
+esp32_ready = DigitalInOut(board.D11)
+esp32_reset = DigitalInOut(board.D12)
+
+spi = busio.SPI(board.SCK, board.MOSI, board.MISO)
+esp = adafruit_esp32spi.ESP_SPIcontrol(spi, esp32_cs, esp32_ready, esp32_reset)
+"""Use below for Most Boards"""
+status_light = neopixel.NeoPixel(
+    board.NEOPIXEL, 1, brightness=0.2
+)  # Uncomment for Most Boards
+"""Uncomment below for ItsyBitsy M4"""
+# status_light = dotstar.DotStar(board.APA102_SCK, board.APA102_MOSI, 1, brightness=0.2)
+# Uncomment below for an externally defined RGB LED
+# import adafruit_rgbled
+# from adafruit_esp32spi import PWMOut
+# RED_LED = PWMOut.PWMOut(esp, 26)
+# GREEN_LED = PWMOut.PWMOut(esp, 27)
+# BLUE_LED = PWMOut.PWMOut(esp, 25)
+# status_light = adafruit_rgbled.RGBLED(RED_LED, BLUE_LED, GREEN_LED)
+wifi = adafruit_esp32spi_wifimanager.ESPSPI_WiFiManager(esp, secrets, status_light)
+
+### Topic Setup ###
+
+# MQTT Topic
+# Use this topic if you'd like to connect to a standard MQTT broker
+mqtt_topic = "test/topic"
+
+# Adafruit IO-style Topic
+# Use this topic if you'd like to connect to io.adafruit.com
+# mqtt_topic = 'aio_user/feeds/temperature'
+
+### Code ###
+
+# Define callback methods which are called when events occur
+# pylint: disable=unused-argument, redefined-outer-name
+def connect(client, userdata, flags, rc):
+    # This function will be called when the client is connected
+    # successfully to the broker.
+    print("Connected to MQTT Broker!")
+    print("Flags: {0}\n RC: {1}".format(flags, rc))
+
+
+def disconnect(client, userdata, rc):
+    # This method is called when the client disconnects
+    # from the broker.
+    print("Disconnected from MQTT Broker!")
+
+
+def subscribe(client, userdata, topic, granted_qos):
+    # This method is called when the client subscribes to a new feed.
+    print("Subscribed to {0} with QOS level {1}".format(topic, granted_qos))
+
+
+def unsubscribe(client, userdata, topic, pid):
+    # This method is called when the client unsubscribes from a feed.
+    print("Unsubscribed from {0} with PID {1}".format(topic, pid))
+
+
+def publish(client, userdata, topic, pid):
+    # This method is called when the client publishes data to a feed.
+    print("Published to {0} with PID {1}".format(topic, pid))
+
+def on_battery_msg(client, topic, message):
+    print("Battery Level: {}v".format(message))
+
+def on_message(client, topic, message):
+    """Method callled when a client's subscribed feed has a new
+    value.
+    :param str topic: The topic of the feed with a new value.
+    :param str message: The new value
+    """
+    print("New message on topic {0}: {1}".format(topic, message))
+
+
+# Connect to WiFi
+print("Connecting to WiFi...")
+wifi.connect()
+print("Connected!")
+
+# Initialize MQTT interface with the esp interface
+MQTT.set_socket(socket, esp)
+
+# Set up a MiniMQTT Client
+client = MQTT.MQTT(broker=secrets["broker"], port=1883)
+
+# Connect callback handlers to client
+client.on_connect = connect
+client.on_disconnect = disconnect
+client.on_subscribe = subscribe
+client.on_unsubscribe = unsubscribe
+client.on_publish = publish
+
+client.on_message = on_message
+client.add_topic_callback("sensors/batteryLevel", on_battery_msg)
+
+# Connect the client to the MQTT broker.
+print("Connecting to MQTT broker...")
+client.connect()
+
+# Subscribe to all notifications on the sensors/
+client.subscribe("sensors/#", 1)
+
+# Start a blocking message loop...
+# NOTE: NO code below this loop will execute
+# NOTE: Network reconnection is handled within this loop
+while True:
+    try:
+        client.loop()
+    except (ValueError, RuntimeError) as e:
+        print("Failed to get data, retrying\n", e)
+        wifi.reset()
+        client.reconnect()
+        continue
+    time.sleep(1)
+
+
