Clean up

bumblebee/modules/caffeine.py

@@ -5,14 +5,16 @@
 Requires the following executables:
     * xdg-screensaver
     * xdotool
+    * xprop (as dependency for xdotool)
     * notify-send
 """
 
+import psutil
+import os
+import logging
 import bumblebee.input
 import bumblebee.output
 import bumblebee.engine
-import psutil
-import os
 
 class Module(bumblebee.engine.Module):
     def __init__(self, engine, config):
@@ -20,37 +22,72 @@ class Module(bumblebee.engine.Module):
             bumblebee.output.Widget(full_text="")
         )
         self._active = False
-        self._xid = 0
+        self._xid = None
         
         engine.input.register_callback(self, button=bumblebee.input.LEFT_MOUSE,
             cmd=self._toggle
         )
 
+    def _check_requirements(self):
+        requirements = ['xdotool','xprop','xdg-screensaver']
+        missing = []
+        for tool in requirements:
+            if not bumblebee.util.which(tool):
+                missing.append(tool)
+        return missing
+
+    def _get_i3bar_xid(self):
+        xid = bumblebee.util.execute("xdotool search --class \"i3bar\"").partition('\n')[0].strip()
+        if xid.isdigit():
+            return xid
+        else:
+            logging.info("Module caffeine: xdotool couldn't get X window ID of \"i3bar\".")
+        return None
+
+    def _notify(self):
+        if not bumblebee.util.which('notify-send'):
+            return
+
+        if self._active:
+            bumblebee.util.execute("notify-send \"Consuming caffeine\"")
+        else:
+            bumblebee.util.execute("notify-send \"Out of coffee\"")
+
+    def _suspend_screensaver(self):
+        self._xid = self._get_i3bar_xid()
+        if self._xid == None:
+            return
+        
+        pid = os.fork()
+        if pid == 0:
+            os.setsid()
+            bumblebee.util.execute("xdg-screensaver suspend {}".format(self._xid))
+            os._exit(0)
+        else:
+            os.waitpid(pid,0)
+                 
+    def _resume_screensaver(self):
+        for process in psutil.process_iter():
+            if process.cmdline() == [bumblebee.util.which('xprop'),'-id',str(self._xid),'-spy']:
+                pid = process.pid
+        os.kill(pid,9)
+
     def state(self, widget):
         if self._active:
             return "activated"
         return "deactivated"
 
     def _toggle(self, event):
+        missing = self._check_requirements()
+        if missing:
+            logging.warning("Could not run caffeine - missing {}!".format(", ".join(missing)))
+            return
+
         self._active = not self._active
-        try:
+        if self._active:
-            if self._active:
+            self._suspend_screensaver()
-                self._xid = bumblebee.util.execute("xdotool search --class \"i3bar\"").strip()
+        else:
-                pid = os.fork()
+            self._resume_screensaver()
-                if pid == 0:
+        self._notify()
-                    os.setsid()
-                    bumblebee.util.execute("xdg-screensaver suspend {}".format(self._xid))
-                    os._exit(0)
-                else:
-                    os.waitpid(pid,0)
-                bumblebee.util.execute("notify-send \"Consuming caffeine\"")
-            else:
-                for process in psutil.process_iter():
-                    if process.cmdline() == ['/usr/bin/xprop','-id',str(self._xid),'-spy']:
-                        pid = process.pid
-                os.kill(pid,9)
-                bumblebee.util.execute("notify-send \"Out of coffee\"")
-        except:
-            self._active = not self._active
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4