From d37ec03dc81117310d377198d1a95d180281cb31 Mon Sep 17 00:00:00 2001
From: Gregory Szorc <gps@mozilla.com>
Date: Wed, 10 Oct 2012 11:08:09 -0700
Subject: [PATCH] Bug 799262 - Formal API for loading mach command modules;
 r=jhammel

---
 mach                     |  1 +
 python/mach/mach/main.py | 82 +++++++++++++++++++++++-----------------
 2 files changed, 49 insertions(+), 34 deletions(-)

diff --git a/mach b/mach
index 743dbbda8c8..bf26540ad00 100755
--- a/mach
+++ b/mach
@@ -44,4 +44,5 @@ except ImportError:
 
 # All of the code is in a module because EVERYTHING IS A LIBRARY.
 mach = mach.main.Mach(our_dir)
+mach.load_commands_from_sys_path()
 sys.exit(mach.run(sys.argv[1:]))
diff --git a/python/mach/mach/main.py b/python/mach/mach/main.py
index 832fa26f4b8..a4fc40198f6 100644
--- a/python/mach/mach/main.py
+++ b/python/mach/mach/main.py
@@ -13,6 +13,7 @@ import imp
 import logging
 import os
 import sys
+import uuid
 
 from mozbuild.base import BuildConfig
 from mozbuild.config import ConfigSettings
@@ -40,8 +41,6 @@ CONSUMED_ARGUMENTS = [
     'func',
 ]
 
-MODULES_SCANNED = False
-
 
 class ArgumentParser(argparse.ArgumentParser):
     """Custom implementation argument parser to make things look pretty."""
@@ -113,10 +112,54 @@ To see more help for a specific command, run:
 
         self.log_manager.register_structured_logger(self.logger)
 
-        if not MODULES_SCANNED:
-            self._load_modules()
+    def load_commands_from_sys_path(self):
+        """Discover and load mach command modules from sys.path.
 
-        MODULES_SCANNED = True
+        This iterates over all paths on sys.path. If the path contains a
+        "mach/commands" subdirectory, all .py files in that directory will be
+        loaded and examined for mach commands.
+        """
+        # Create parent module otherwise Python complains when the parent is
+        # missing.
+        if b'mach.commands' not in sys.modules:
+            mod = imp.new_module(b'mach.commands')
+            sys.modules[b'mach.commands'] = mod
+
+        for path in sys.path:
+            # We only support importing .py files from directories.
+            commands_path = os.path.join(path, 'mach', 'commands')
+
+            if not os.path.isdir(commands_path):
+                continue
+
+            self.load_commands_from_directory(commands_path)
+
+    def load_commands_from_directory(self, path):
+        """Scan for mach commands from modules in a directory.
+
+        This takes a path to a directory, loads the .py files in it, and
+        registers and found mach command providers with this mach instance.
+        """
+        for f in sorted(os.listdir(path)):
+            if not f.endswith('.py') or f == '__init__.py':
+                continue
+
+            full_path = os.path.join(path, f)
+            module_name = 'mach.commands.%s' % f[0:-3]
+
+            self.load_commands_from_file(full_path, module_name=module_name)
+
+    def load_commands_from_file(self, path, module_name=None):
+        """Scan for mach commands from a file.
+
+        This takes a path to a file and loads it as a Python module under the
+        module name specified. If no name is specified, a random one will be
+        chosen.
+        """
+        if module_name is None:
+            module_name = 'mach.commands.%s' % uuid.uuid1().get_hex()
+
+        imp.load_source(module_name, path)
 
     def run(self, argv):
         """Runs mach with arguments provided from the command line.
@@ -217,35 +260,6 @@ To see more help for a specific command, run:
         self.logger.log(level, format_str,
             extra={'action': action, 'params': params})
 
-    def _load_modules(self):
-        """Scan over Python modules looking for mach command providers."""
-
-        # Create parent module otherwise Python complains when the parent is
-        # missing.
-        if b'mach.commands' not in sys.modules:
-            mod = imp.new_module(b'mach.commands')
-            sys.modules[b'mach.commands'] = mod
-
-        for path in sys.path:
-            # We only support importing .py files from directories.
-            commands_path = os.path.join(path, 'mach', 'commands')
-
-            if not os.path.isdir(commands_path):
-                continue
-
-            # We only support loading modules in the immediate mach.commands
-            # module, not sub-modules. Walking the tree would be trivial to
-            # implement if it were ever desired.
-            for f in sorted(os.listdir(commands_path)):
-                if not f.endswith('.py') or f == '__init__.py':
-                    continue
-
-                full_path = os.path.join(commands_path, f)
-                module_name = 'mach.commands.%s' % f[0:-3]
-
-                imp.load_source(module_name, full_path)
-
-
     def load_settings(self, args):
         """Determine which settings files apply and load them.