Add PSUControl support

octoprint_homeassistant/__init__.py

@@ -53,6 +53,7 @@ def __init__(self):
         self.mqtt_subcribe = None
         self.update_timer = None
         self.constant_timer = None
+        self.psucontrol_enabled = False
 
     def handle_timer(self):
         self._generate_printer_status()
@@ -121,6 +122,46 @@ def on_after_startup(self):
                     self._on_mqtt_message,
                 )
 
+        # PSUControl helpers
+        psu_helpers = self._plugin_manager.get_helpers(
+            "psucontrol", "turn_psu_on", "turn_psu_off", "get_psu_state"
+        )
+
+        self.psucontrol_enabled = True
+
+        if psu_helpers:
+            self._logger.info("PSUControl helpers found")
+
+            if "get_psu_state" in psu_helpers:
+                self.get_psu_state = psu_helpers["get_psu_state"]
+                self._logger.debug("Setup get_psu_state helper")
+            else:
+                self._logger.error(
+                    "Helper get_psu_state not found, disabling PSUControl integration"
+                )
+                self.psucontrol_enabled = False
+
+            if "turn_psu_on" in psu_helpers:
+                self.turn_psu_on = psu_helpers["turn_psu_on"]
+                self._logger.debug("Setup turn_psu_on helper")
+            else:
+                self._logger.error(
+                    "Helper turn_psu_on not found, disabling PSUControl integration"
+                )
+                self.psucontrol_enabled = False
+
+            if "turn_psu_off" in psu_helpers:
+                self.turn_psu_off = psu_helpers["turn_psu_off"]
+                self._logger.debug("Setup turn_psu_off helper")
+            else:
+                self._logger.error(
+                    "Helper turn_psu_on not found, disabling PSUControl integration"
+                )
+                self.psucontrol_enabled = False
+        else:
+            self._logger.info("PSUControl helpers not found")
+            self.psucontrol_enabled = False
+
         self.snapshot_enabled = self._settings.global_get(
             ["webcam", "timelapseEnabled"]
         )
@@ -160,6 +201,9 @@ def on_after_startup(self):
         self.on_print_progress("", "", 0)
         self._generate_connection_status()
 
+        if self.psucontrol_enabled:
+            self._generate_psu_state()
+
     def _get_mac_address(self):
         import uuid
 
@@ -581,6 +625,20 @@ def _generate_connection_status(self):
                 allow_queueing=True,
             )
 
+    def _generate_psu_state(self, psu_state=None):
+        if self.psucontrol_enabled:
+            if psu_state is None:
+                psu_state = self.get_psu_state()
+                self._logger.debug(
+                    "No psu_state specified, state retrieved from helper: "
+                    + str(psu_state)
+                )
+            self.mqtt_publish(
+                self._generate_topic("hassTopic", "psu_on", full=True),
+                str(psu_state),
+                allow_queueing=True,
+            )
+
     def _on_emergency_stop(
         self, topic, message, retained=None, qos=None, *args, **kwargs
     ):
@@ -615,6 +673,16 @@ def _on_shutdown_system(self, topic, message, retained=None, qos=None, *args, **
             except Exception as e:
                 self._logger.info("Unable to run shutdown command: " + str(e))
 
+    def _on_psu(self, topic, message, retained=None, qos=None, *args, **kwargs):
+        message = message.decode()
+        self._logger.debug("PSUControl message received: " + message)
+        if message == "True":
+            self._logger.info("Turning on PSU")
+            self.turn_psu_on()
+        else:
+            self._logger.info("Turning off PSU")
+            self.turn_psu_off()
+
     def _on_camera(self, topic, message, retained=None, qos=None, *args, **kwargs):
         self._logger.debug("Camera snapshot message received: " + str(message))
         if self.snapshot_enabled:
@@ -798,6 +866,28 @@ def _generate_device_controls(self, subscribe=False):
             },
         )
 
+        # PSUControl
+        if self.psucontrol_enabled:
+            if subscribe:
+                self.mqtt_subscribe(
+                    self._generate_topic("controlTopic", "psu", full=True),
+                    self._on_psu,
+                )
+
+            self._generate_sensor(
+                topic=_discovery_topic + "/switch/" + _node_id + "_PSU/config",
+                values={
+                    "name": _node_name + " PSU",
+                    "uniq_id": _node_id + "_PSU",
+                    "cmd_t": "~" + self._generate_topic("controlTopic", "psu"),
+                    "stat_t": "~" + self._generate_topic("hassTopic", "psu_on"),
+                    "pl_on": "True",
+                    "pl_off": "False",
+                    "device": _config_device,
+                    "ic": "mdi:flash",
+                },
+            )
+
         # Camera output
         if self.snapshot_enabled:
             if subscribe:
@@ -936,6 +1026,13 @@ def on_event(self, event, payload):
                 allow_queueing=True,
             )
 
+
+        if (
+            self.psucontrol_enabled and 
+            event == Events.PLUGIN_PSUCONTROL_PSU_STATE_CHANGED
+        ):
+            self._generate_psu_state(payload["isPSUOn"])
+
         if event == Events.CAPTURE_DONE:
             file_handle = open(payload["file"], "rb")
             file_content = file_handle.read()