Add color and UI overrides to the distutils frontend, and clean up the nocolor stuff a tad.

coderanger · coderanger · commit 45859f8bb0d2 · 2010-08-09T05:46:39.000-07:00
diff --git a/pyzen/distutils.py b/pyzen/distutils.py
@@ -6,35 +6,35 @@
 from pkg_resources import *
 
 from pyzen.core import main
-from pyzen.runner import ColoredTextTestRunner
+from pyzen.runner import get_test_runner
 
 class ZenTestProgram(unittest.TestProgram):
     
     def runTests(self):
-        if self.testRunner is None:
-            self.testRunner = unittest.TextTestRunner
-
-        if isinstance(self.testRunner, (type, types.ClassType)):
-            try:
-                testRunner = self.testRunner(verbosity=self.verbosity)
-            except TypeError:
-                # didn't accept the verbosity argument
-                testRunner = self.testRunner()
-        else:
-            # it is assumed to be a TestRunner instance
-            testRunner = self.testRunner
+        testRunner = self.testRunner(verbosity=self.verbosity)
         self.result = testRunner.run(self.test)
 
-def run_tests(loader, args, path):
+def run_tests(loader, args, path, nocolor):
     sys.path[:] = path
     loader_ep = EntryPoint.parse("x="+loader)
     loader_class = loader_ep.load(require=False)
-    m = ZenTestProgram(None, None, [unittest.__file__]+args, testLoader=loader_class())
+    m = ZenTestProgram(None, None, [unittest.__file__]+args, testLoader=loader_class(), testRunner=get_test_runner(nocolor))
     return m.result
 
 class zen(test):
     """Command to run test suite under PyZen."""
     description = 'run unit tests under PyZen'
-
+    
+    user_options = test.user_options + [
+        ('ui=', 'u', 'Force the use of the given PyZen UI'),
+        ('nocolor', 'c', 'Disable colored output'),
+    ]
+    boolean_options = ['nocolor']
+    
+    def initialize_options(self):
+        test.initialize_options(self)
+        self.ui = None
+        self.nocolor = False
+    
     def run_tests(self):
-        main(None, run_tests, self.test_loader, self.test_args, sys.path[:])
+        main(self.ui, run_tests, self.test_loader, self.test_args, sys.path[:], self.nocolor)
diff --git a/pyzen/flaskext/script.py b/pyzen/flaskext/script.py
@@ -7,7 +7,7 @@
 from flaskext.script import Command, Option
 
 from pyzen.core import main
-from pyzen.runner import ColoredTextTestRunner
+from pyzen.runner import get_test_runner
 
 try:
     from unittest2 import TestLoader
@@ -31,11 +31,7 @@ def run_tests(app, pattern, start_dir, verbosity, nocolor):
     print start_dir
     loader = ZenTestLoader()
     suite = loader.discover(start_dir, pattern, start_dir)
-    if nocolor:
-        runner = unittest.TextTestRunner
-    else:
-        runner = ColoredTextTestRunner
-    result = runner(verbosity=verbosity).run(suite)
+    result = get_test_runner(nocolor)(verbosity=verbosity).run(suite)
     return result
 
 class Test(Command):
diff --git a/pyzen/management/commands/zen.py b/pyzen/management/commands/zen.py
@@ -12,7 +12,7 @@
     patch_for_test_db_setup = lambda: None
     
 from pyzen.core import main
-from pyzen.runner import ColoredTextTestRunner
+from pyzen.runner import get_test_runner
 
 def run_tests(*test_labels, **options):
     patch_for_test_db_setup()
@@ -24,10 +24,7 @@ def run_tests(*test_labels, **options):
     
     class NewTestSuiteRunner(TestSuiteRunner):
         def run_suite(self, suite, **kwargs):
-            if nocolor:
-                return unittest.TextTestRunner(verbosity=self.verbosity).run(suite)
-            else:
-                return ColoredTextTestRunner(verbosity=self.verbosity).run(suite)
+            return get_test_runner(nocolor)(verbosity=self.verbosity).run(suite)
         def suite_result(self, suite, result, **kwargs):
             return result
     
diff --git a/pyzen/runner.py b/pyzen/runner.py
@@ -71,3 +71,9 @@ def run(self, test):
         else:
             self.stream.writeln(COLOR_SUCCESS+"OK"+COLOR_RESET)
         return result
+
+def get_test_runner(nocolor):
+    if nocolor:
+        return unittest.TextTestRunner
+    else:
+        return ColoredTextTestRunner
