move win color utils in _common.py

2 files changed, 35 insertions, 32 deletions

diff --git a/psutil/_common.py b/psutil/_common.py
index 1f8900df..5b2d4214 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -9,7 +9,9 @@
 
 from __future__ import division, print_function
 
+import atexit
 import contextlib
+import ctypes
 import errno
 import functools
 import os
@@ -789,6 +791,37 @@ def hilite(s, color="green", bold=False):
     return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
 
 
+def print_color(s, color="green", bold=False, file=sys.stdout):
+    def _stderr_handle():
+        GetStdHandle = ctypes.windll.Kernel32.GetStdHandle
+        STD_ERROR_HANDLE_ID = ctypes.c_ulong(0xfffffff4)
+        GetStdHandle.restype = ctypes.c_ulong
+        handle = GetStdHandle(STD_ERROR_HANDLE_ID)
+        atexit.register(ctypes.windll.Kernel32.CloseHandle, handle)
+        return handle
+
+    def win_colorprint(s, color, bold=False):
+        DEFAULT_COLOR = 7
+        colors = dict(green=2, red=4, brown=6)
+        try:
+            color = colors[color]
+        except KeyError:
+            raise ValueError("invalid color %r; choose between %r" % (
+                color, list(colors.keys())))
+        if bold and color <= 7:
+            color += 8
+        handle = _stderr_handle()
+        SetConsoleTextAttribute = \
+            ctypes.windll.Kernel32.SetConsoleTextAttribute
+        SetConsoleTextAttribute(handle, color)
+        try:
+            print(s)
+        finally:
+            SetConsoleTextAttribute(handle, DEFAULT_COLOR)
+
+    win_colorprint(s, color, bold)
+
+
 if bool(os.getenv('PSUTIL_DEBUG', 0)):
     import inspect
 
diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
index bfa9c5fc..f4e9e8b1 100755
--- a/psutil/tests/runner.py
+++ b/psutil/tests/runner.py
@@ -10,7 +10,6 @@ on KeyboardInterrupt.
 """
 
 from __future__ import print_function
-import atexit
 import os
 import sys
 import unittest
@@ -24,7 +23,7 @@ except ImportError:
 
 import psutil
 from psutil._common import hilite
-from psutil._common import memoize
+from psutil._common import print_color
 from psutil._common import term_supports_colors
 from psutil.tests import safe_rmpath
 from psutil.tests import TOX
@@ -35,42 +34,13 @@ VERBOSITY = 1 if TOX else 2
 FAILED_TESTS_FNAME = '.failed-tests.txt'
 
 
-@memoize
-def _stderr_handle():
-    GetStdHandle = ctypes.windll.Kernel32.GetStdHandle
-    STD_ERROR_HANDLE_ID = ctypes.c_ulong(0xfffffff4)
-    GetStdHandle.restype = ctypes.c_ulong
-    handle = GetStdHandle(STD_ERROR_HANDLE_ID)
-    atexit.register(ctypes.windll.Kernel32.CloseHandle, handle)
-    return handle
-
-
-def win_colorprint(printer, s, color, bold=False):
-    DEFAULT_COLOR = 7
-    colors = dict(green=2, red=4, brown=6)
-    try:
-        color = colors[color]
-    except KeyError:
-        raise ValueError("invalid color %r; choose between %r" % (
-            list(colors.keys())))
-    if bold and color <= 7:
-        color += 8
-    handle = _stderr_handle()
-    SetConsoleTextAttribute = ctypes.windll.Kernel32.SetConsoleTextAttribute
-    SetConsoleTextAttribute(handle, color)
-    try:
-        printer(s)
-    finally:
-        SetConsoleTextAttribute(handle, DEFAULT_COLOR)
-
-
 class ColouredResult(TextTestResult):
 
     def _color_print(self, s, color, bold=False):
         if os.name == 'posix':
             self.stream.writeln(hilite(s, color, bold=bold))
         else:
-            win_colorprint(self.stream.writeln, s, color, bold=bold)
+            print_color(s, color, bold=bold)
 
     def addSuccess(self, test):
         TestResult.addSuccess(self, test)