Autodetect modules

.travis.yml

@@ -27,12 +27,13 @@ before_install:
  - python -V
  - pip -V
  - pip3 install --upgrade pip
- - pip3 install pyinstaller
+ - pip3 install mypy pyinstaller
 
 install:
  - make build
 
 script:
+ - make typecheck
  - make test
  - make test-integration
  - make package

aw_qt/manager.py

@@ -1,65 +1,119 @@
 import os
 import platform
+from glob import glob
 from time import sleep
 import logging
 import subprocess
-from typing import Optional, List
+import shutil
+from typing import Optional, List, Dict
 
 import aw_core
 
 logger = logging.getLogger(__name__)
 
+_module_dir = os.path.dirname(os.path.realpath(__file__))
+_parent_dir = os.path.abspath(os.path.join(_module_dir, os.pardir))
+_search_paths = [_module_dir, _parent_dir]
 
-def _locate_executable(name: str) -> List[str]:
-    """
-    Will start module from localdir if present there,
-    otherwise will try to call what is available in PATH.
 
-    Returns it as a Popen cmd list.
-    """
-    curr_filepath = os.path.realpath(__file__)
-    curr_dir = os.path.dirname(curr_filepath)
-    search_paths = [curr_dir, os.path.abspath(os.path.join(curr_dir, os.pardir))]
-    exec_paths = [os.path.join(path, name) for path in search_paths]
+def _locate_bundled_executable(name: str) -> Optional[str]:
+    """Returns the path to the module executable if it exists in the bundle, else None."""
+    _exec_paths = [os.path.join(path, name) for path in _search_paths]
 
-    for exec_path in exec_paths:
+    # Look for it in the installation path
+    for exec_path in _exec_paths:
         if os.path.isfile(exec_path):
             # logger.debug("Found executable for {} in: {}".format(name, exec_path))
-            return [exec_path]
-            break  # this break is redundant, but kept due to for-else semantics
+            return exec_path
+    return None
+
+
+def _is_system_module(name) -> bool:
+    """Checks if a module with a particular name exists in PATH"""
+    return shutil.which(name) is not None
+
+
+def _locate_executable(name: str) -> Optional[str]:
+    """
+    Will return the path to the executable if bundled,
+    otherwise returns the name if it is available in PATH.
+
+    Used when calling Popen.
+    """
+    exec_path = _locate_bundled_executable(name)
+    if exec_path is not None:  # Check if it exists in bundle
+        return exec_path
+    elif _is_system_module(name):  # Check if it's in PATH
+        return name
     else:
-        # TODO: Actually check if it is in PATH
-        # logger.debug("Trying to start {} using PATH (executable not found in: {})"
-        #              .format(name, exec_paths))
-        return [name]
+        logger.warning("Could not find module '{}' in installation directory or PATH".format(name))
+        return None
+
+
+def _discover_modules_bundled() -> List[str]:
+    # Look for modules in source dir and parent dir
+    modules = []
+    for path in _search_paths:
+        matches = glob(os.path.join(path, "aw-*"))
+        for match in matches:
+            if os.path.isfile(match) and os.access(match, os.X_OK):
+                name = os.path.basename(match)
+                modules.append(name)
+            else:
+                logger.warning("Found matching file but was not executable: {}".format(path))
+
+    logger.info("Found bundled modules: {}".format(set(modules)))
+    return modules
+
+
+def _discover_modules_system() -> List[str]:
+    search_paths = os.environ["PATH"].split(":")
+    modules = []
+    for path in search_paths:
+        if os.path.isdir(path):
+            files = os.listdir(path)
+            for filename in files:
+                if "aw-" in filename:
+                    modules.append(filename)
+
+    logger.info("Found system modules: {}".format(set(modules)))
+    return modules
 
 
 class Module:
     def __init__(self, name: str, testing: bool = False) -> None:
         self.name = name
-        self.started = False
+        self.started = False  # Should be True if module is supposed to be running, else False
         self.testing = testing
+        self.location = "system" if _is_system_module(name) else "bundled"
         self._process = None  # type: Optional[subprocess.Popen]
         self._last_process = None  # type: Optional[subprocess.Popen]
 
     def start(self) -> None:
         logger.info("Starting module {}".format(self.name))
 
         # Create a process group, become its leader
+        # TODO: This shouldn't go here
         if platform.system() != "Windows":
             os.setpgrp()
 
-        exec_cmd = _locate_executable(self.name)
-        if self.testing:
-            exec_cmd.append("--testing")
-        # logger.debug("Running: {}".format(exec_cmd))
+        exec_path = _locate_executable(self.name)
+        if exec_path is None:
+            return
+        else:
+            exec_cmd = [exec_path]
+            if self.testing:
+                exec_cmd.append("--testing")
+            # logger.debug("Running: {}".format(exec_cmd))
 
-        # There is a very good reason stdout and stderr is not PIPE here
-        # See: https://github.com/ActivityWatch/aw-server/issues/27
-        self._process = subprocess.Popen(exec_cmd, universal_newlines=True)
+            # There is a very good reason stdout and stderr is not PIPE here
+            # See: https://github.com/ActivityWatch/aw-server/issues/27
+            try:
+                self._process = subprocess.Popen(exec_cmd, universal_newlines=True)
+            except OSError as e:
+                logger.error("Couldn't start module with command {} ({})".format(exec_cmd, e))
 
-        # Should be True if module is supposed to be running, else False
-        self.started = True
+            self.started = True
 
     def stop(self) -> None:
         """
@@ -108,19 +162,21 @@ def read_log(self) -> str:
 
 
 class Manager:
-    def __init__(self, testing: bool=False):
-        # TODO: Fetch these from somewhere appropriate (auto detect or a config file)
-        #       Save to config wether they should autostart or not.
-        _possible_modules = [
-            "aw-server",
-            "aw-watcher-afk",
-            "aw-watcher-window",
-            # "aw-watcher-spotify",
-            # "aw-watcher-network"
-        ]
-
-        # TODO: Filter away all modules not available on system
-        self.modules = {name: Module(name, testing=testing) for name in _possible_modules}
+    def __init__(self, testing: bool = False) -> None:
+        self.testing = testing
+        self.modules = {}  # type: Dict[str, Module]
+
+        self.discover_modules()
+
+    def discover_modules(self):
+        # These should always be bundled with aw-qt
+        found_modules = set(_discover_modules_bundled())
+        found_modules |= set(_discover_modules_system())
+        found_modules ^= {"aw-qt"}  # Exclude self
+
+        for m_name in found_modules:
+            if m_name not in self.modules:
+                self.modules[m_name] = Module(m_name, testing=self.testing)
 
     def get_unexpected_stops(self):
         return list(filter(lambda x: x.started and not x.is_alive(), self.modules.values()))

aw_qt/trayicon.py

@@ -4,6 +4,7 @@
 import webbrowser
 import os
 import subprocess
+from collections import defaultdict
 
 from PyQt5 import QtCore
 from PyQt5.QtWidgets import QApplication, QSystemTrayIcon, QMessageBox, QMenu, QWidget, QPushButton
@@ -37,7 +38,7 @@ def open_dir(d):
 
 
 class TrayIcon(QSystemTrayIcon):
-    def __init__(self, manager: Manager, icon, parent=None, testing=False):
+    def __init__(self, manager: Manager, icon, parent=None, testing=False) -> None:
         QSystemTrayIcon.__init__(self, icon, parent)
         self.setToolTip("ActivityWatch" + (" (testing)" if testing else ""))
 
@@ -93,12 +94,18 @@ def show_module_failed_dialog(module):
             box.show()
 
         def rebuild_modules_menu():
-            for module in modulesMenu.actions():
-                name = module.text()
-                alive = self.manager.modules[name].is_alive()
-                module.setChecked(alive)
-                # print(module.text(), alive)
+            for action in modulesMenu.actions():
+                if action.isEnabled():
+                    name = action.module.name
+                    alive = self.manager.modules[name].is_alive()
+                    action.setChecked(alive)
+                    # print(module.text(), alive)
 
+            # TODO: Do it in a better way, singleShot isn't pretty...
+            QtCore.QTimer.singleShot(2000, rebuild_modules_menu)
+        QtCore.QTimer.singleShot(2000, rebuild_modules_menu)
+
+        def check_module_status():
             unexpected_exits = self.manager.get_unexpected_stops()
             if unexpected_exits:
                 for module in unexpected_exits:
@@ -107,22 +114,34 @@ def rebuild_modules_menu():
 
             # TODO: Do it in a better way, singleShot isn't pretty...
             QtCore.QTimer.singleShot(2000, rebuild_modules_menu)
-
-        QtCore.QTimer.singleShot(2000, rebuild_modules_menu)
+        QtCore.QTimer.singleShot(2000, check_module_status)
 
     def _build_modulemenu(self, moduleMenu):
         moduleMenu.clear()
 
         def add_module_menuitem(module):
-            ac = moduleMenu.addAction(module.name, lambda: module.toggle())
+            title = module.name
+            ac = moduleMenu.addAction(title, lambda: module.toggle())
+            # Kind of nasty, but a quick way to affiliate the module to the menu action for when it needs updating
+            ac.module = module
             ac.setCheckable(True)
             ac.setChecked(module.is_alive())
 
-        add_module_menuitem(self.manager.modules["aw-server"])
+        modules_by_location = defaultdict(lambda: list())
+        for module in sorted(self.manager.modules.values(), key=lambda m: m.name):
+            modules_by_location[module.location].append(module)
+
+        for location, modules in sorted(modules_by_location.items(), key=lambda kv: kv[0]):
+            header = moduleMenu.addAction(location)
+            header.setEnabled(False)
+
+            # Always show aw-server first
+            if "aw-server" in (m.name for m in modules):
+                add_module_menuitem(self.manager.modules["aw-server"])
 
-        for module_name in sorted(self.manager.modules.keys()):
-            if module_name != "aw-server":
-                add_module_menuitem(self.manager.modules[module_name])
+            for module in sorted(modules, key=lambda m: m.name):
+                if module.name != "aw-server":
+                    add_module_menuitem(self.manager.modules[module.name])
 
 
 def exit_dialog():
@@ -177,7 +196,3 @@ def run(manager, testing=False):
 
     # Run the application, blocks until quit
     return app.exec_()
-
-
-if __name__ == "__main__":
-    run()