Replace MQTT functions with class

source/main.py

@@ -5,13 +5,14 @@
 from snmp.v1.exceptions import NoSuchName
 
 from functions import *
-from mqtt_publisher import *
+from mqtt_publisher import MQTTPublisher
 from settings import *
 from snmp_channels import CHANNELS
 
 FAIL_COUNT = 0
 
-publish_meta(DEVICE_NAME, '')
+mqtt = MQTTPublisher(MQTT_BROKER_IP, ROOT_MQTT_TOPIC)
+mqtt.publish_meta(DEVICE_NAME)
 
 while True:
     manager = Manager(SNMP_COMMUNITY)
@@ -30,23 +31,23 @@
                         mqtt_type = table[name]['Type']
                         mqtt_units = table[name]['Units']
                         mqtt_order = table[name]['Order']
-                        publish_control(data=mqtt_data,
-                                        name=name,
-                                        data_type=mqtt_type,
-                                        units=mqtt_units,
-                                        error='',
-                                        order=mqtt_order)
+                        mqtt.publish_control(data=mqtt_data,
+                                             name=name,
+                                             data_type=mqtt_type,
+                                             units=mqtt_units,
+                                             error='',
+                                             order=mqtt_order)
                 else:
                     mqtt_type = CHANNELS[key]['Type']
                     mqtt_units = CHANNELS[key]['Units']
                     response = manager.get(host, oid)
                     value = transform_item(response[0], key)
-                    publish_control(data=value,
-                                    name=key,
-                                    data_type=mqtt_type,
-                                    units=mqtt_units,
-                                    error='',
-                                    order=mqtt_order)
+                    mqtt.publish_control(data=value,
+                                         name=key,
+                                         data_type=mqtt_type,
+                                         units=mqtt_units,
+                                         error='',
+                                         order=mqtt_order)
 
             except NoSuchName:
                 continue
@@ -61,6 +62,6 @@
         manager.close()
 
         if FAIL_COUNT > 2:
-            publish_error('Work_status')
+            mqtt.publish_error('Work_status')
 
         time.sleep(POLLING_INTERVAL)

source/mqtt_publisher.py

@@ -1,58 +1,53 @@
-from paho.mqtt.publish import single, multiple
-
-from settings import *
-
-
-def publish_meta(name: str, error: str):
-    msgs = [
-        {'topic': '{}/meta/name'.format(ROOT_MQTT_TOPIC),
-         'payload': name,
-         'retain': True},
-
-        {'topic': '{}/meta/error'.format(ROOT_MQTT_TOPIC),
-         'payload': error,
-         'retain': True},
-    ]
-
-    multiple(msgs, hostname=MQTT_BROKER_IP)
-
-
-def publish_control(data,
-                    name: str,
-                    data_type: str,
-                    units: str,
-                    error: str,
-                    order=None,
-                    retain=True):
-    msgs = [
-        {'topic': '{}/controls/{}'.format(ROOT_MQTT_TOPIC, name),
-         'payload': data,
-         'retain': retain},
-
-        {'topic': '{}/controls/{}/meta/type'.format(ROOT_MQTT_TOPIC, name),
-         'payload': data_type,
-         'retain': retain},
-
-        {'topic': '{}/controls/{}/meta/units'.format(ROOT_MQTT_TOPIC, name),
-         'payload': units,
-         'retain': retain},
-
-        {'topic': '{}/controls/{}/meta/order'.format(ROOT_MQTT_TOPIC, name),
-         'payload': order,
-         'retain': retain},
-
-        {'topic': '{}/controls/{}/meta/error'.format(ROOT_MQTT_TOPIC, name),
-         'payload': error,
-         'retain': retain},
-    ]
-
-    multiple(msgs, hostname=MQTT_BROKER_IP)
-
-
-def publish_error(name: str, retain: bool = True) -> None:
-    single(
-        topic='{}/controls/{}/meta/error'.format(ROOT_MQTT_TOPIC, name),
-        payload='r',
-        retain=retain,
-        hostname=MQTT_BROKER_IP,
-    )
+import paho.mqtt.client as mqtt
+
+
+class MQTTPublisher:
+    def __init__(self, broker_ip: str, root_topic: str):
+        self.client = mqtt.Client()
+        self.broker_ip = broker_ip
+        self.root_topic = root_topic
+        self.client.connect(self.broker_ip)
+
+    def __del__(self):
+        self.disconnect()
+
+    def publish_meta(self, name: str, error: str = '') -> None:
+        msgs = [
+            ('{}/meta/name'.format(self.root_topic), name, 1, True),
+            ('{}/meta/error'.format(self.root_topic), error, 1, True),
+        ]
+        self.publish_multiple(msgs)
+
+    def publish_control(
+            self,
+            data,
+            name: str,
+            data_type: str,
+            units: str,
+            error: str,
+            order=None,
+            retain=True,
+    ) -> None:
+        msgs = [
+            ('{}/controls/{}'.format(self.root_topic, name), data, 1, retain),
+            ('{}/controls/{}/meta/type'.format(self.root_topic, name), data_type, 1, retain),
+            ('{}/controls/{}/meta/units'.format(self.root_topic, name), units, 1, retain),
+            ('{}/controls/{}/meta/order'.format(self.root_topic, name), order, 1, retain),
+            ('{}/controls/{}/meta/error'.format(self.root_topic, name), error, 1, retain),
+        ]
+        self.publish_multiple(msgs)
+
+    def publish_error(self, name: str, retain: bool = True) -> None:
+        self.client.publish(
+            '{}/controls/{}/meta/error'.format(self.root_topic, name),
+            'r',
+            retain=retain
+        )
+
+    def publish_multiple(self, msgs) -> None:
+        for msg in msgs:
+            topic, payload, qos, retain = msg
+            self.client.publish(topic, payload, qos, retain)
+
+    def disconnect(self) -> None:
+        self.client.disconnect()