refactor print colors utils

10 files changed, 133 insertions, 208 deletions

diff --git a/MANIFEST.in b/MANIFEST.in
index 038d4baa..3b4232e9 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,3 +1,4 @@
+include .cirrus.yml
 include .coveragerc
 include .gitignore
 include CREDITS
@@ -114,7 +115,6 @@ include scripts/internal/README
 include scripts/internal/bench_oneshot.py
 include scripts/internal/bench_oneshot_2.py
 include scripts/internal/check_broken_links.py
-include scripts/internal/download_exes.py
 include scripts/internal/fix_flake8.py
 include scripts/internal/generate_manifest.py
 include scripts/internal/print_access_denied.py
@@ -122,7 +122,7 @@ include scripts/internal/print_announce.py
 include scripts/internal/print_api_speed.py
 include scripts/internal/print_timeline.py
 include scripts/internal/purge_installation.py
-include scripts/internal/scriptutils.py
+include scripts/internal/win_download_wheels.py
 include scripts/internal/winmake.py
 include scripts/iotop.py
 include scripts/killall.py
diff --git a/Makefile b/Makefile
index 890c6e41..812ba9c2 100644
--- a/Makefile
+++ b/Makefile
@@ -11,20 +11,24 @@ DEPS = \
 	check-manifest \
 	coverage \
 	flake8 \
-	futures \
-	ipaddress \
-	mock==1.0.1 \
 	pyperf \
 	requests \
 	setuptools \
-	sphinx \
 	twine \
-	unittest2 \
 	virtualenv \
 	wheel
 
+ifeq ($(PYTHON), $(filter $(PYTHON), python python2 python2.7))
+	DEPS += \
+		futures \
+		ipaddress \
+		mock==1.0.1 \
+		unittest2
+endif
+
 # In not in a virtualenv, add --user options for install commands.
-INSTALL_OPTS = `$(PYTHON) -c "import sys; print('' if hasattr(sys, 'real_prefix') else '--user')"`
+INSTALL_OPTS = `$(PYTHON) -c \
+	"import sys; print('' if hasattr(sys, 'real_prefix') else '--user')"`
 TEST_PREFIX = PYTHONWARNINGS=all PSUTIL_TESTING=1 PSUTIL_DEBUG=1
 
 all: test
@@ -197,7 +201,7 @@ wheel:  ## Generate wheel.
 	$(PYTHON) setup.py bdist_wheel
 
 win-download-wheels:  ## Download wheels hosted on appveyor.
-	$(TEST_PREFIX) $(PYTHON) scripts/internal/download_exes.py --user giampaolo --project psutil
+	$(TEST_PREFIX) $(PYTHON) scripts/internal/win_download_wheels.py --user giampaolo --project psutil
 
 upload-src:  ## Upload source tarball on https://pypi.org/project/psutil/
 	${MAKE} sdist
diff --git a/docs/Makefile b/docs/Makefile
index c7f4723a..cca5435f 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -15,6 +15,9 @@ ALLSPHINXOPTS   = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
 # the i18n builder cannot share the environment and doctrees with the others
 I18NSPHINXOPTS  = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
 
+DEPS = sphinx
+
+
 .PHONY: help
 help:
 	@echo "Please use \`make <target>' where <target> is one of"
@@ -224,3 +227,7 @@ dummy:
 	$(SPHINXBUILD) -b dummy $(ALLSPHINXOPTS) $(BUILDDIR)/dummy
 	@echo
 	@echo "Build finished. Dummy builder generates no files."
+
+.PHONY: setup-dev-env
+setup-dev-env:  ## Install GIT hooks, pip, test deps (also upgrades them).
+	$(PYTHON) -m pip install --user --upgrade --trusted-host files.pythonhosted.org $(DEPS)
diff --git a/psutil/_common.py b/psutil/_common.py
index 9306cd15..728d9c62 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -43,10 +43,9 @@ else:
 PY3 = sys.version_info[0] == 3
 
 __all__ = [
-    # constants
+    # OS constants
     'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'MACOS', 'OSX', 'POSIX',
     'SUNOS', 'WINDOWS',
-    'ENCODING', 'ENCODING_ERRS', 'AF_INET6',
     # connection constants
     'CONN_CLOSE', 'CONN_CLOSE_WAIT', 'CONN_CLOSING', 'CONN_ESTABLISHED',
     'CONN_FIN_WAIT1', 'CONN_FIN_WAIT2', 'CONN_LAST_ACK', 'CONN_LISTEN',
@@ -58,6 +57,8 @@ __all__ = [
     'STATUS_RUNNING', 'STATUS_SLEEPING', 'STATUS_STOPPED', 'STATUS_SUSPENDED',
     'STATUS_TRACING_STOP', 'STATUS_WAITING', 'STATUS_WAKE_KILL',
     'STATUS_WAKING', 'STATUS_ZOMBIE', 'STATUS_PARKED',
+    # other constants
+    'ENCODING', 'ENCODING_ERRS', 'AF_INET6',
     # named tuples
     'pconn', 'pcputimes', 'pctxsw', 'pgids', 'pio', 'pionice', 'popenfile',
     'pthread', 'puids', 'sconn', 'scpustats', 'sdiskio', 'sdiskpart',
@@ -66,7 +67,9 @@ __all__ = [
     'conn_tmap', 'deprecated_method', 'isfile_strict', 'memoize',
     'parse_environ_block', 'path_exists_strict', 'usage_percent',
     'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum', "wrap_numbers",
-    'bytes2human', 'conn_to_ntuple', 'hilite', 'debug',
+    'bytes2human', 'conn_to_ntuple', 'debug',
+    # shell utils
+    'hilite', 'term_supports_colors', 'print_color',
 ]
 
 
@@ -757,38 +760,78 @@ else:
         return s
 
 
-def _term_supports_colors(file=sys.stdout):
-    if hasattr(_term_supports_colors, "ret"):
-        return _term_supports_colors.ret
+# =====================================================================
+# --- shell utils
+# =====================================================================
+
+
+@memoize
+def term_supports_colors(file=sys.stdout):
+    if os.name == 'nt':
+        return True
     try:
         import curses
         assert file.isatty()
         curses.setupterm()
         assert curses.tigetnum("colors") > 0
     except Exception:
-        _term_supports_colors.ret = False
         return False
     else:
-        _term_supports_colors.ret = True
-    return _term_supports_colors.ret
+        return True
 
 
-def hilite(s, ok=True, bold=False):
+def hilite(s, color="green", bold=False):
     """Return an highlighted version of 'string'."""
-    if not _term_supports_colors():
+    if not term_supports_colors():
         return s
     attr = []
-    if ok is None:  # no color
-        pass
-    elif ok:   # green
-        attr.append('32')
-    else:   # red
-        attr.append('31')
+    colors = dict(green='32', red='91', brown='33')
+    colors[None] = '29'
+    try:
+        color = colors[color]
+    except KeyError:
+        raise ValueError("invalid color %r; choose between %r" % (
+            list(colors.keys())))
+    attr.append(color)
     if bold:
         attr.append('1')
     return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
 
 
+def print_color(s, color="green", bold=False, file=sys.stdout):
+    """Print a colorized version of string."""
+    if not term_supports_colors():
+        print(s, file=file)
+    elif POSIX:
+        print(hilite(s, color, bold), file=file)
+    else:
+        import ctypes
+
+        DEFAULT_COLOR = 7
+        GetStdHandle = ctypes.windll.Kernel32.GetStdHandle
+        SetConsoleTextAttribute = \
+            ctypes.windll.Kernel32.SetConsoleTextAttribute
+
+        colors = dict(green=2, red=4, brown=6)
+        colors[None] = DEFAULT_COLOR
+        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_id = -12 if file is sys.stderr else -11
+        GetStdHandle.restype = ctypes.c_ulong
+        handle = GetStdHandle(handle_id)
+        SetConsoleTextAttribute(handle, color)
+        try:
+            print(s, file=file)
+        finally:
+            SetConsoleTextAttribute(handle, DEFAULT_COLOR)
+
+
 if bool(os.getenv('PSUTIL_DEBUG', 0)):
     import inspect
 
diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
index f8601bad..4c3359dd 100755
--- a/psutil/tests/runner.py
+++ b/psutil/tests/runner.py
@@ -5,12 +5,13 @@
 # found in the LICENSE file.
 
 """
-Unit test runner, providing colourized output and printing failures
-on KeyboardInterrupt.
+Unit test runner, providing new features on top of unittest module:
+- colourized output (error, skip)
+- print failures/tracebacks on CTRL+C
+- re-run failed tests only (make test-failed)
 """
 
 from __future__ import print_function
-import atexit
 import os
 import sys
 import unittest
@@ -23,7 +24,9 @@ except ImportError:
     ctypes = None
 
 import psutil
-from psutil._common import memoize
+from psutil._common import hilite
+from psutil._common import print_color
+from psutil._common import term_supports_colors
 from psutil.tests import safe_rmpath
 from psutil.tests import TOX
 
@@ -31,95 +34,37 @@ from psutil.tests import TOX
 HERE = os.path.abspath(os.path.dirname(__file__))
 VERBOSITY = 1 if TOX else 2
 FAILED_TESTS_FNAME = '.failed-tests.txt'
-if os.name == 'posix':
-    GREEN = 1
-    RED = 2
-    BROWN = 94
-else:
-    GREEN = 2
-    RED = 4
-    BROWN = 6
-    DEFAULT_COLOR = 7
-
-
-def term_supports_colors(file=sys.stdout):
-    if os.name == 'nt':
-        return ctypes is not None
-    try:
-        import curses
-        assert file.isatty()
-        curses.setupterm()
-        assert curses.tigetnum("colors") > 0
-    except Exception:
-        return False
-    else:
-        return True
-
-
-def hilite(s, color, bold=False):
-    """Return an highlighted version of 'string'."""
-    attr = []
-    if color == GREEN:
-        attr.append('32')
-    elif color == RED:
-        attr.append('91')
-    elif color == BROWN:
-        attr.append('33')
-    else:
-        raise ValueError("unrecognized color")
-    if bold:
-        attr.append('1')
-    return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
-
-
-@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):
-    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)
+
+
+# =====================================================================
+# --- unittest subclasses
+# =====================================================================
 
 
 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)
+    def _print_color(self, s, color, bold=False):
+        file = sys.stderr if color == "red" else sys.stdout
+        print_color(s, color, bold=bold, file=file)
 
     def addSuccess(self, test):
         TestResult.addSuccess(self, test)
-        self._color_print("OK", GREEN)
+        self._print_color("OK", "green")
 
     def addError(self, test, err):
         TestResult.addError(self, test, err)
-        self._color_print("ERROR", RED, bold=True)
+        self._print_color("ERROR", "red", bold=True)
 
     def addFailure(self, test, err):
         TestResult.addFailure(self, test, err)
-        self._color_print("FAIL", RED)
+        self._print_color("FAIL", "red")
 
     def addSkip(self, test, reason):
         TestResult.addSkip(self, test, reason)
-        self._color_print("skipped: %s" % reason, BROWN)
+        self._print_color("skipped: %s" % reason, "brown")
 
     def printErrorList(self, flavour, errors):
-        flavour = hilite(flavour, RED, bold=flavour == 'ERROR')
+        flavour = hilite(flavour, "red", bold=flavour == 'ERROR')
         TextTestResult.printErrorList(self, flavour, errors)
 
 
@@ -133,6 +78,11 @@ class ColouredRunner(TextTestRunner):
         return self.result
 
 
+# =====================================================================
+# --- public API
+# =====================================================================
+
+
 def setup_tests():
     if 'PSUTIL_TESTING' not in os.environ:
         # This won't work on Windows but set_testing() below will do it.
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 0277a56a..e5ff6e45 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -311,7 +311,6 @@ class TestProcess(unittest.TestCase):
     @skip_on_not_implemented(only_if=LINUX)
     def test_io_counters(self):
         p = psutil.Process()
-
         # test reads
         io1 = p.io_counters()
         with open(PYTHON_EXE, 'rb') as f:
diff --git a/scripts/internal/print_access_denied.py b/scripts/internal/print_access_denied.py
index 2c757fd7..9123ba6d 100644
--- a/scripts/internal/print_access_denied.py
+++ b/scripts/internal/print_access_denied.py
@@ -50,7 +50,7 @@ from collections import defaultdict
 import time
 
 import psutil
-from scriptutils import hilite
+from psutil._common import print_color
 
 
 def main():
@@ -75,13 +75,12 @@ def main():
     # print
     templ = "%-20s %-5s %-9s %s"
     s = templ % ("API", "AD", "Percent", "Outcome")
-    print(hilite(s, ok=None, bold=True))
+    print_color(s, color=None, bold=True)
     for methname, ads in sorted(d.items(), key=lambda x: (x[1], x[0])):
         perc = (ads / tot_procs) * 100
         outcome = "SUCCESS" if not ads else "ACCESS DENIED"
         s = templ % (methname, ads, "%6.1f%%" % perc, outcome)
-        s = hilite(s, ok=not ads)
-        print(s)
+        print_color(s, "red" if ads else None)
     tot_perc = round((tot_ads / tot_calls) * 100, 1)
     print("-" * 50)
     print("Totals: access-denied=%s (%s%%), calls=%s, processes=%s, "
diff --git a/scripts/internal/print_api_speed.py b/scripts/internal/print_api_speed.py
index a99293c4..85d1cfc5 100644
--- a/scripts/internal/print_api_speed.py
+++ b/scripts/internal/print_api_speed.py
@@ -9,53 +9,45 @@
 $ make print_api_speed
 SYSTEM APIS               SECONDS
 ----------------------------------
-boot_time                 0.000140
-cpu_count                 0.000016
-cpu_count (cores)         0.000312
-cpu_freq                  0.000811
-cpu_percent               0.000138
-cpu_stats                 0.000165
-cpu_times                 0.000140
+cpu_count                 0.000014
+disk_usage                0.000027
+cpu_times                 0.000037
+cpu_percent               0.000045
 ...
 
 PROCESS APIS              SECONDS
 ----------------------------------
-children                  0.007246
-cmdline                   0.000069
-connections               0.000072
-cpu_affinity              0.000012
-cpu_num                   0.000035
-cpu_percent               0.000042
-cpu_times                 0.000031
+create_time               0.000001
+nice                      0.000005
+cwd                       0.000011
+cpu_affinity              0.000011
+ionice                    0.000013
+...
 """
 
 from __future__ import print_function, division
 from timeit import default_timer as timer
-import argparse
 import inspect
 import os
 
 import psutil
-from scriptutils import hilite
+from psutil._common import print_color
 
 
-SORT_BY_TIME = False if psutil.POSIX else True
-TOP_SLOWEST = 7
 timings = []
 templ = "%-25s %s"
 
 
 def print_timings():
-    slower = []
-    timings.sort(key=lambda x: x[1 if SORT_BY_TIME else 0])
-    for x in sorted(timings, key=lambda x: x[1], reverse=1)[:TOP_SLOWEST]:
-        slower.append(x[0])
+    timings.sort(key=lambda x: x[1])
+    i = 0
     while timings[:]:
         title, elapsed = timings.pop(0)
         s = templ % (title, "%f" % elapsed)
-        if title in slower:
-            s = hilite(s, ok=False)
-        print(s)
+        if i > len(timings) - 5:
+            print_color(s, color="red")
+        else:
+            print(s)
 
 
 def timecall(title, fun, *args, **kw):
@@ -65,11 +57,7 @@ def timecall(title, fun, *args, **kw):
     timings.append((title, elapsed))
 
 
-def titlestr(s):
-    return hilite(s, ok=None, bold=True)
-
-
-def run():
+def main():
     # --- system
 
     public_apis = []
@@ -83,7 +71,7 @@ def run():
             if name not in ignore:
                 public_apis.append(name)
 
-    print(titlestr(templ % ("SYSTEM APIS", "SECONDS")))
+    print_color(templ % ("SYSTEM APIS", "SECONDS"), color=None, bold=True)
     print("-" * 34)
     for name in public_apis:
         fun = getattr(psutil, name)
@@ -99,7 +87,7 @@ def run():
 
     # --- process
     print("")
-    print(titlestr(templ % ("PROCESS APIS", "SECONDS")))
+    print_color(templ % ("PROCESS APIS", "SECONDS"), color=None, bold=True)
     print("-" * 34)
     ignore = ['send_signal', 'suspend', 'resume', 'terminate', 'kill', 'wait',
               'as_dict', 'parent', 'parents', 'memory_info_ex', 'oneshot',
@@ -114,18 +102,5 @@ def run():
     print_timings()
 
 
-def main():
-    global SORT_BY_TIME, TOP_SLOWEST
-    parser = argparse.ArgumentParser(description='Benchmark all API calls')
-    parser.add_argument('-t', '--time', required=False, default=SORT_BY_TIME,
-                        action='store_true', help="sort by timings")
-    parser.add_argument('-s', '--slowest', required=False, default=TOP_SLOWEST,
-                        help="highlight the top N slowest APIs")
-    args = parser.parse_args()
-    SORT_BY_TIME = bool(args.time)
-    TOP_SLOWEST = int(args.slowest)
-    run()
-
-
 if __name__ == '__main__':
     main()
diff --git a/scripts/internal/scriptutils.py b/scripts/internal/scriptutils.py
deleted file mode 100644
index 3ee416f8..00000000
--- a/scripts/internal/scriptutils.py
+++ /dev/null
@@ -1,54 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-"""Utils shared by all files in scripts/internal."""
-
-from __future__ import print_function
-import sys
-
-from psutil._compat import lru_cache
-
-
-__all__ = ['hilite', 'printerr', 'exit']
-
-
-@lru_cache()
-def _term_supports_colors(file=sys.stdout):
-    try:
-        import curses
-        assert file.isatty()
-        curses.setupterm()
-        assert curses.tigetnum("colors") > 0
-    except Exception:
-        return False
-    else:
-        return True
-
-
-def hilite(s, ok=True, bold=False):
-    """Return an highlighted version of 'string'."""
-    if not _term_supports_colors():
-        return s
-    attr = []
-    if ok is None:  # no color
-        pass
-    elif ok:   # green
-        attr.append('32')
-    else:   # red
-        attr.append('31')
-    if bold:
-        attr.append('1')
-    return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
-
-
-def printerr(s):
-    print(hilite(s, ok=False), file=sys.stderr)
-
-
-def exit(msg=""):
-    if msg:
-        printerr(msg)
-    sys.exit(1)
diff --git a/scripts/internal/download_exes.py b/scripts/internal/win_download_wheels.py
index 4a559bb0..0cb37afe 100755
--- a/scripts/internal/download_exes.py
+++ b/scripts/internal/win_download_wheels.py
@@ -19,10 +19,11 @@ import errno
 import os
 import requests
 import shutil
+import sys
 
 from psutil import __version__ as PSUTIL_VERSION
 from psutil._common import bytes2human
-from scriptutils import printerr, exit
+from psutil._common import print_color
 
 
 BASE_URL = 'https://ci.appveyor.com/api'
@@ -81,7 +82,8 @@ def get_file_urls(options):
                 file_url = job_url + '/' + item['fileName']
                 urls.append(file_url)
         if not urls:
-            exit("no artifacts found")
+            print_color("no artifacts found", 'ret')
+            sys.exit(1)
         else:
             for url in sorted(urls, key=lambda x: os.path.basename(x)):
                 yield url
@@ -111,7 +113,7 @@ def run(options):
             try:
                 local_fname = fut.result()
             except Exception:
-                printerr("error while downloading %s" % (url))
+                print_color("error while downloading %s" % (url), 'red')
                 raise
             else:
                 completed += 1