refactor print colors utils

4 files changed, 28 insertions, 106 deletions

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