Implemented simple color output and moved __main__ to main().

__init__.py

@@ -1,2 +1,2 @@
-from runner import TextTestRunner
-__all__ = ['TextTestRunner']
+from runner import TextTestRunner, main
+__all__ = ['TextTestRunner', 'main']

__main__.py

@@ -6,39 +6,15 @@
 #
 
 import sys
-import os
 import unittest
 
-from testrunner import TextTestRunner
+from testrunner import TextTestRunner, main
 
 if len(sys.argv) < 2:
     print 'Usage: %s tests/test_COMPONENT.py' % sys.argv[0]
     sys.exit(1)
 
-
 # TODO parse verbosity option '-v', given in argv.
-verbosity = 0
-
-suites = []
-
-# Dynamic load the requested module containing the test cases.
-for testfile in sys.argv[1:]:
-
-    module_dir, module_file = os.path.split(testfile)
-    module_name, module_ext = os.path.splitext(module_file)
-
-    module_obj = __import__(module_dir.replace('/', '.'),
-                            fromlist=[module_name])
-    module_obj.__file__ = testfile
-    globals()[module_name] = module_obj
-
-    # Start the test runner and display the results to stdout.
-    container = module_obj.__dict__[module_name]
-    class_parts = module_name[5:].split('_')
-    class_name = 'Test' + ''.join([p.capitalize() for p in class_parts])
-
-    testcase = container.__dict__[class_name]
-    suites += [unittest.TestLoader().loadTestsFromTestCase(testcase)]
+verbose = 0
 
-# Create the text runner and execute the tests.
-TextTestRunner(verbosity=verbosity).run(unittest.TestSuite(suites))
+main(sys.argv[1:], verbose)

runner.py

@@ -1,5 +1,7 @@
-import unittest
+import os
+import sys
 import time
+import unittest
 
 
 class TextTestRunner(unittest.TextTestRunner):
@@ -7,6 +9,11 @@ class TextTestRunner(unittest.TextTestRunner):
     printed to stdout. The method ``run`` is modified and it's original source
     is in Python Standard Library's ``unittest`` module."""
 
+    def __init__(self, stream=sys.stderr, descriptions=True, verbosity=1,
+            color=True):
+        unittest.TextTestRunner.__init__(self, stream, descriptions, verbosity)
+        self.color = color
+
     def run(self, test):
         "Run the given test case or test suite."
         result = self._makeResult()
@@ -29,8 +36,8 @@ class TextTestRunner(unittest.TextTestRunner):
         result.printErrors()
 
         run = result.testsRun
-        self.stream.write("Ran %d test%s in %.3fs " %
-                          (run, run != 1 and "s" or "", timeTaken))
+        msg = 'Ran %d test%s in %.3fs '
+        self.stream.write(msg % (run, run != 1 and 's' or '', timeTaken))
 
         expectedFails = unexpectedSuccesses = skipped = 0
         try:
@@ -45,25 +52,75 @@ class TextTestRunner(unittest.TextTestRunner):
         infos = []
 
         if not result.wasSuccessful():
-            self.stream.write("FAILED")
+            if self.color:
+                msg = '\033[0;31mFAILED\033[1;m'
+            else:
+                msg = 'FAILED'
+            self.stream.write(msg)
+
             failed, errored = map(len, (result.failures, result.errors))
-            if failed:
-                infos.append("failures=%d" % failed)
+            if failed: 
+                if self.color:
+                    msg = 'failures=\033[1;31m%d\033[1;m'
+                else:
+                    msg = 'failures=%d'
+                infos.append(msg % failed)
             if errored:
-                infos.append("errors=%d" % errored)
+                if self.color:
+                    msg = 'errors=\033[1;31m%d\033[1;m'
+                else:
+                    msg = 'errors=%d'
+                infos.append(msg % errored)
         else:
-            self.stream.write("OK")
+            if self.color:
+                msg = '\033[0;32mOK\033[0;m'
+            else:
+                msg = 'OK'
+            self.stream.write(msg)
 
         if skipped:
-            infos.append("skipped=%d" % skipped)
+            if self.color:
+                msg = 'skipped=\033[1;33m%d\033[1;m'
+            else:
+                msg = 'skipped=%d'
+            infos.append(msg % skipped)
+
         if expectedFails:
-            infos.append("expected failures=%d" % expectedFails)
+            infos.append('expected failures=%d' % expectedFails)
+
         if unexpectedSuccesses:
-            infos.append("unexpected successes=%d" % unexpectedSuccesses)
+            infos.append('unexpected successes=%d' % unexpectedSuccesses)
 
         if infos:
-            self.stream.writeln(" (%s)" % (", ".join(infos),))
+            self.stream.writeln(' (%s)' % (', '.join(infos),))
         else:
-            self.stream.write("\n")
+            self.stream.write('\n')
 
         return result
+
+def main(tests, verbose=0, color=True):
+    suites = []
+
+    # Dynamic load the requested module containing the test cases.
+    for testfile in tests:
+
+        module_dir, module_file = os.path.split(testfile)
+        module_name, module_ext = os.path.splitext(module_file)
+
+        # XXX: this looks really hacky, so cleanup is necessary.
+        module_obj = __import__(module_dir.replace('/', '.'),
+                                fromlist=[module_name])
+        module_obj.__file__ = testfile
+        globals()[module_name] = module_obj
+
+        # Start the test runner and display the results to stdout.
+        container = module_obj.__dict__[module_name]
+        class_parts = module_name[5:].split('_')
+        class_name = 'Test' + ''.join([p.capitalize() for p in class_parts])
+
+        testcase = container.__dict__[class_name]
+        suites += [unittest.TestLoader().loadTestsFromTestCase(testcase)]
+
+    # Create the text runner and execute the tests.
+    runner = TextTestRunner(verbosity=verbose, color=color)
+    runner.run(unittest.TestSuite(suites))