Formatting cleanups for UPS and Prometheus functionality

bmspy/prometheus.py

@@ -1,11 +1,15 @@
+import time
 import prometheus_client
 from bmspy.utilities import debugger
+from bmspy.server import collect_data
 
 
 def prometheus_export(daemonize=True, filename=None):
     global debug
     if not can_export_prometheus:
-        raise ModuleNotFoundError("Unable to export to Prometheus. Is prometheus-client installed?")
+        raise ModuleNotFoundError(
+            "Unable to export to Prometheus. Is prometheus-client installed?"
+        )
 
     data = dict()
     # Initialize data structure, to fill in help values
@@ -34,44 +38,50 @@ def prometheus_export(daemonize=True, filename=None):
             prometheus_client.generate_latest(registry)
     else:
         if filename is None:
-            debugger("Invalid filename supplied");
+            debugger("Invalid filename supplied")
             return False
         prometheus_client.write_to_textfile(filename, registry=registry)
         return True
 
+
 def prometheus_create_metric(registry, data):
     metric = dict()
     for name, contains in data.items():
-        helpmsg = ''
+        helpmsg = ""
-        if contains.get('help') is not None:
+        if contains.get("help") is not None:
-            helpmsg = contains.get('help')
+            helpmsg = contains.get("help")
-            if contains.get('units'):
+            if contains.get("units"):
-                helpmsg += ' (' + contains.get('units') + ')'
+                helpmsg += " (" + contains.get("units") + ")"
-        if contains.get('value') is not None:
+        if contains.get("value") is not None:
             metric[name] = prometheus_client.Gauge(name, helpmsg, registry=registry)
         # Has multiple values, each a different label
-        elif contains.get('values') is not None:
+        elif contains.get("values") is not None:
-            if contains.get('label') is None:
+            if contains.get("label") is None:
                 debugger("ERROR: no label for {0} specified".format(name))
-            label = contains.get('label')
+            label = contains.get("label")
-            metric[name] = prometheus_client.Gauge(name, helpmsg, [label], registry=registry)
+            metric[name] = prometheus_client.Gauge(
-        elif contains.get('info') is not None:
+                name, helpmsg, [label], registry=registry
+            )
+        elif contains.get("info") is not None:
             metric[name] = prometheus_client.Info(name, helpmsg, registry=registry)
         else:
             pass
     return metric
 
+
 def prometheus_populate_metric(metric, data):
     for name, contains in data.items():
-        if contains.get('value') is not None:
+        if contains.get("value") is not None:
-            value = contains.get('value')
+            value = contains.get("value")
             metric[name].set(value)
         # doesn't have a value, but has [1-4]:
-        if contains.get('values') is not None and isinstance(contains.get('values'), dict):
+        if contains.get("values") is not None and isinstance(
-            for idx, label_value in contains.get('values').items():
+            contains.get("values"), dict
+        ):
+            for idx, label_value in contains.get("values").items():
                 metric[name].labels(idx).set(label_value)
-        if contains.get('info'):
+        if contains.get("info"):
-            value = contains.get('info')
+            value = contains.get("info")
             metric[name].info({name: value})
         else:
             pass
@@ -79,9 +89,11 @@ def prometheus_populate_metric(metric, data):
 
 # TODO fork bms daemon if need be?
 
+
 def main():
     debugger("TODO. At present, run from bmspy directly.")
 
+
 #    influxdb_export(bucket=args.influx_bucket, \
 #                    url=args.influx_url, \
 #                    org=args.influx_org, \

bmspy/ups.py

@@ -1,8 +1,13 @@
-from collections import deque
 import argparse
-import atexit, datetime, os, re, sys, time
+import atexit
-import smtplib, ssl, socket
+import os
-from typing import Any
+import re
+import time
+import smtplib
+import ssl
+import socket
+from collections import deque
+
 from bmspy import client
 
 DAEMON_UPDATE_PERIOD = 30
@@ -12,21 +17,31 @@ critical_sent = False
 warning_sent = False
 alert_sent = False
 
-def handle_shutdown(action: str = 'cancel', delay: int = 0, debug: int = 0) -> None:
+
+def handle_shutdown(action: str = "cancel", delay: int = 0, debug: int = 0) -> None:
     global scheduled_shutdown
 
-    if action == 'shutdown':
+    if action == "shutdown":
         if scheduled_shutdown is False:
             scheduled_shutdown = time.time() + delay * 60 * 1000
             os.system("/sbin/shutdown {}".format(delay))
 
-    elif action == 'cancel':
+    elif action == "cancel":
         os.system("/sbin/shutdown -c")
 
     return
 
 
-def handle_email(text: str, level: str | None, recipient: str = "root", mailserver: str = "localhost", port: int = 25, mailuser: str | None = None, mailpass: str | None = None, debug: int = 0) -> None:
+def handle_email(
+    text: str,
+    level: str | None,
+    recipient: str = "root",
+    mailserver: str = "localhost",
+    port: int = 25,
+    mailuser: str | None = None,
+    mailpass: str | None = None,
+    debug: int = 0,
+) -> None:
     isSSL = False
     hostname = socket.gethostname()
 
@@ -38,7 +53,9 @@ def handle_email(text: str, level: str | None, recipient: str = "root", mailserv
         isSSL = True
 
     if level is not None:
-        msg = "From: {}\r\nTo: {}\r\nSubject: {} from BMSPY UPS on {}\r\n\r\n{}\r\n".format(sender, recipient, level, hostname, text)
+        msg = "From: {}\r\nTo: {}\r\nSubject: {} from BMSPY UPS on {}\r\n\r\n{}\r\n".format(
+            sender, recipient, level, hostname, text
+        )
 
     if isSSL:
         context = ssl.create_default_context()
@@ -58,43 +75,105 @@ def main() -> None:
     global alert_sent, warning_sent, critical_sent
 
     parser = argparse.ArgumentParser(
-                description='Query JBD BMS and alert or shutdown when certain thresholds are reached',
+        description="Query JBD BMS and alert or shutdown when certain thresholds are reached",
         add_help=True,
     )
-    parser.add_argument('--alert', '-a', dest='alert', action='store_true',
+    parser.add_argument(
-                default=True, help='Email an alert when UPS detects A/C loss (default: true)')
+        "--alert",
-    parser.add_argument('--warning', '-w', dest='warning_threshold', action='store', type=int,
+        "-a",
-                default=75, help='Email an alert when remaining capacity percentage drops below this figure (default: 75)')
+        dest="alert",
-    parser.add_argument('--critical', '-c', dest='critical_threshold', action='store', type=int,
+        action="store_true",
-                default=30, help='Shut system down when remaining capacity percentage drops below this figure (default: 30)')
+        default=True,
-    parser.add_argument('--delay', '-d', dest='shutdown_delay', action='store', type=int,
+        help="Email an alert when UPS detects A/C loss (default: true)",
-                default=5, help='Delay system shutdown (default: 5 minutes)')
+    )
-    parser.add_argument('--mailserver', '-m', dest='mailserver', action='store',
+    parser.add_argument(
-                default="localhost", help='Mail server (default: localhost)')
+        "--warning",
-    parser.add_argument('--port', '-p', dest='port', action='store',
+        "-w",
-                default=25, help='Mail server port (default: 25)')
+        dest="warning_threshold",
-    parser.add_argument('--user', dest='mailuser', action='store',
+        action="store",
-                default=None, help='Mail server user')
+        type=int,
-    parser.add_argument('--pass', dest='mailpass', action='store',
+        default=75,
-                default=None, help='Mail server password')
+        help="Email an alert when remaining capacity percentage drops below this figure (default: 75)",
-    parser.add_argument('--to', '-t', dest='recipient', action='store',
+    )
-                default="root", help='Email recipient (default: root)')
+    parser.add_argument(
-    parser.add_argument('--socket', '-s', dest='socket', action='store',
+        "--critical",
-                default='/run/bmspy/bms', help='Socket to communicate with daemon')
+        "-c",
-    parser.add_argument('--verbose', '-v', action='count',
+        dest="critical_threshold",
-                default=0, help='Print more verbose information (can be specified multiple times)')
+        action="store",
+        type=int,
+        default=30,
+        help="Shut system down when remaining capacity percentage drops below this figure (default: 30)",
+    )
+    parser.add_argument(
+        "--delay",
+        "-d",
+        dest="shutdown_delay",
+        action="store",
+        type=int,
+        default=5,
+        help="Delay system shutdown (default: 5 minutes)",
+    )
+    parser.add_argument(
+        "--mailserver",
+        "-m",
+        dest="mailserver",
+        action="store",
+        default="localhost",
+        help="Mail server (default: localhost)",
+    )
+    parser.add_argument(
+        "--port",
+        "-p",
+        dest="port",
+        action="store",
+        default=25,
+        help="Mail server port (default: 25)",
+    )
+    parser.add_argument(
+        "--user", dest="mailuser", action="store", default=None, help="Mail server user"
+    )
+    parser.add_argument(
+        "--pass",
+        dest="mailpass",
+        action="store",
+        default=None,
+        help="Mail server password",
+    )
+    parser.add_argument(
+        "--to",
+        "-t",
+        dest="recipient",
+        action="store",
+        default="root",
+        help="Email recipient (default: root)",
+    )
+    parser.add_argument(
+        "--socket",
+        "-s",
+        dest="socket",
+        action="store",
+        default="/run/bmspy/bms",
+        help="Socket to communicate with daemon",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="count",
+        default=0,
+        help="Print more verbose information (can be specified multiple times)",
+    )
     args = parser.parse_args()
 
     debug = args.verbose
 
     print("Running BMS UPS daemon on socket {}".format(args.socket))
 
-    client.handle_registration(args.socket, 'ups', debug)
+    client.handle_registration(args.socket, "ups", debug)
-    atexit.register(client.handle_registration, args.socket, 'ups', debug)
+    atexit.register(client.handle_registration, args.socket, "ups", debug)
 
     history = deque()
     while True:
-        data = client.read_data(args.socket, 'ups')
+        data = client.read_data(args.socket, "ups")
         history.append(data)
 
         # Remove the oldest data from the history
@@ -111,79 +190,111 @@ def main() -> None:
             time.sleep(DAEMON_UPDATE_PERIOD)
             continue
 
-        current_amps = float(data['bms_current_amps']['raw_value'])
+        current_amps = float(data["bms_current_amps"]["raw_value"])
-        charge_ratio = float(data['bms_capacity_charge_ratio']['raw_value']) * 100
+        charge_ratio = float(data["bms_capacity_charge_ratio"]["raw_value"]) * 100
-        comparison_1_current_amps = float(comparison_1['bms_current_amps']['raw_value'])
+        comparison_1_current_amps = float(comparison_1["bms_current_amps"]["raw_value"])
-        comparison_1_charge_ratio = float(comparison_1['bms_capacity_charge_ratio']['raw_value']) * 100
+        comparison_1_charge_ratio = (
-        comparison_2_current_amps = float(comparison_2['bms_current_amps']['raw_value'])
+            float(comparison_1["bms_capacity_charge_ratio"]["raw_value"]) * 100
-        comparison_2_charge_ratio = float(comparison_2['bms_capacity_charge_ratio']['raw_value']) * 100
+        )
-        comparison_3_current_amps = float(comparison_3['bms_current_amps']['raw_value'])
+        comparison_2_current_amps = float(comparison_2["bms_current_amps"]["raw_value"])
-        comparison_3_charge_ratio = float(comparison_3['bms_capacity_charge_ratio']['raw_value']) * 100
+        comparison_2_charge_ratio = (
+            float(comparison_2["bms_capacity_charge_ratio"]["raw_value"]) * 100
+        )
+        comparison_3_current_amps = float(comparison_3["bms_current_amps"]["raw_value"])
+        comparison_3_charge_ratio = (
+            float(comparison_3["bms_capacity_charge_ratio"]["raw_value"]) * 100
+        )
 
         if debug > 1:
-            print("current:            {:>3.2f}A\ncapacity remaining: {:>4.0f}%".format(current_amps, charge_ratio))
+            print(
+                "current:            {:>3.2f}A\ncapacity remaining: {:>4.0f}%".format(
+                    current_amps, charge_ratio
+                )
+            )
 
-        if charge_ratio <= args.critical_threshold and \
+        if (
-           comparison_1_charge_ratio <= args.critical_threshold:
+            charge_ratio <= args.critical_threshold
+            and comparison_1_charge_ratio <= args.critical_threshold
+        ):
             if debug > 0:
                 print("Below critical threshold, shutting down")
-            handle_shutdown(action = 'shutdown', delay = args.shutdown_delay, debug = debug)
+            handle_shutdown(action="shutdown", delay=args.shutdown_delay, debug=debug)
             if critical_sent is False:
-                handle_email(text = "remaining capacity below {}%, shutting down".format(args.critical_threshold),
+                handle_email(
+                    text="remaining capacity below {}%, shutting down".format(
+                        args.critical_threshold
+                    ),
                     level="Critical alert",
                     recipient=args.recipient,
                     mailserver=args.mailserver,
                     port=args.port,
                     mailuser=args.mailuser,
                     mailpass=args.mailpass,
-                    debug = debug)
+                    debug=debug,
+                )
                 critical_sent = True
 
-        elif charge_ratio <= args.warning_threshold and \
+        elif (
-           comparison_1_charge_ratio <= args.warning_threshold:
+            charge_ratio <= args.warning_threshold
+            and comparison_1_charge_ratio <= args.warning_threshold
+        ):
             if debug > 0:
                 print("Below warning threshold")
             if warning_sent is False:
-                handle_email(text = "remaining capacity below {}%".format(args.warning_threshold),
+                handle_email(
+                    text="remaining capacity below {}%".format(args.warning_threshold),
                     level="Warning",
                     recipient=args.recipient,
                     mailserver=args.mailserver,
                     port=args.port,
                     mailuser=args.mailuser,
                     mailpass=args.mailpass,
-                    debug = debug)
+                    debug=debug,
+                )
                 warning_sent = True
 
         # Current needs to be negative for two consecutive reads
-        elif args.alert and current_amps < 0 and \
+        elif (
-          comparison_1_current_amps < 0 and \
+            args.alert
-          comparison_2_current_amps >= 0:
+            and current_amps < 0
+            and comparison_1_current_amps < 0
+            and comparison_2_current_amps >= 0
+        ):
             if debug > 0:
                 print("Alert: discharging!")
             if alert_sent is False:
-                handle_email(text = "power lost", level = "Power loss alert",
+                handle_email(
+                    text="power lost",
+                    level="Power loss alert",
                     recipient=args.recipient,
                     mailserver=args.mailserver,
                     port=args.port,
                     mailuser=args.mailuser,
                     mailpass=args.mailpass,
-                    debug = debug)
+                    debug=debug,
+                )
                 alert_sent = True
 
         # Current needs to be zero or positive for two consecutive reads
-        elif args.alert and current_amps >= 0 and \
+        elif (
-          comparison_1_current_amps >= 0 and \
+            args.alert
-          comparison_2_current_amps < 0:
+            and current_amps >= 0
+            and comparison_1_current_amps >= 0
+            and comparison_2_current_amps < 0
+        ):
             if debug > 0:
                 print("Alert: power regained!")
-            handle_shutdown(action = 'cancel', debug = debug)
+            handle_shutdown(action="cancel", debug=debug)
-            handle_email(text = "power regained", level = "Recovery alert",
+            handle_email(
+                text="power regained",
+                level="Recovery alert",
                 recipient=args.recipient,
                 mailserver=args.mailserver,
                 port=args.port,
                 mailuser=args.mailuser,
                 mailpass=args.mailpass,
-                debug = debug)
+                debug=debug,
+            )
             critical_sent = False
             warning_sent = False
             alert_sent = False