1#!/usr/bin/env python3
2
3# Copyright (c) 2009 Giampaolo Rodola'. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7
8"""Shortcuts for various tasks, emulating UNIX "make" on Windows.
9This is supposed to be invoked by "make.bat" and not used directly.
10This was originally written as a bat file but they suck so much
11that they should be deemed illegal!
12"""
13
14from __future__ import print_function
15import argparse
16import atexit
17import ctypes
18import errno
19import fnmatch
20import os
21import shutil
22import site
23import ssl
24import subprocess
25import sys
26import tempfile
27
28
29APPVEYOR = bool(os.environ.get('APPVEYOR'))
30if APPVEYOR:
31    PYTHON = sys.executable
32else:
33    PYTHON = os.getenv('PYTHON', sys.executable)
34RUNNER_PY = 'psutil\\tests\\runner.py'
35GET_PIP_URL = "https://bootstrap.pypa.io/get-pip.py"
36PY3 = sys.version_info[0] == 3
37HERE = os.path.abspath(os.path.dirname(__file__))
38ROOT_DIR = os.path.realpath(os.path.join(HERE, "..", ".."))
39PYPY = '__pypy__' in sys.builtin_module_names
40DEPS = [
41    "coverage",
42    "flake8",
43    "nose",
44    "pdbpp",
45    "pip",
46    "pyperf",
47    "pyreadline",
48    "setuptools",
49    "wheel",
50    "requests"
51]
52if sys.version_info[:2] <= (2, 7):
53    DEPS.append('unittest2')
54if sys.version_info[:2] <= (2, 7):
55    DEPS.append('mock')
56if sys.version_info[:2] <= (3, 2):
57    DEPS.append('ipaddress')
58if sys.version_info[:2] <= (3, 4):
59    DEPS.append('enum34')
60if not PYPY:
61    DEPS.append("pywin32")
62    DEPS.append("wmi")
63
64_cmds = {}
65if PY3:
66    basestring = str
67
68GREEN = 2
69LIGHTBLUE = 3
70YELLOW = 6
71RED = 4
72DEFAULT_COLOR = 7
73
74
75# ===================================================================
76# utils
77# ===================================================================
78
79
80def safe_print(text, file=sys.stdout, flush=False):
81    """Prints a (unicode) string to the console, encoded depending on
82    the stdout/file encoding (eg. cp437 on Windows). This is to avoid
83    encoding errors in case of funky path names.
84    Works with Python 2 and 3.
85    """
86    if not isinstance(text, basestring):
87        return print(text, file=file)
88    try:
89        file.write(text)
90    except UnicodeEncodeError:
91        bytes_string = text.encode(file.encoding, 'backslashreplace')
92        if hasattr(file, 'buffer'):
93            file.buffer.write(bytes_string)
94        else:
95            text = bytes_string.decode(file.encoding, 'strict')
96            file.write(text)
97    file.write("\n")
98
99
100def stderr_handle():
101    GetStdHandle = ctypes.windll.Kernel32.GetStdHandle
102    STD_ERROR_HANDLE_ID = ctypes.c_ulong(0xfffffff4)
103    GetStdHandle.restype = ctypes.c_ulong
104    handle = GetStdHandle(STD_ERROR_HANDLE_ID)
105    atexit.register(ctypes.windll.Kernel32.CloseHandle, handle)
106    return handle
107
108
109def win_colorprint(s, color=LIGHTBLUE):
110    color += 8  # bold
111    handle = stderr_handle()
112    SetConsoleTextAttribute = ctypes.windll.Kernel32.SetConsoleTextAttribute
113    SetConsoleTextAttribute(handle, color)
114    try:
115        print(s)
116    finally:
117        SetConsoleTextAttribute(handle, DEFAULT_COLOR)
118
119
120def sh(cmd, nolog=False):
121    if not nolog:
122        safe_print("cmd: " + cmd)
123    p = subprocess.Popen(cmd, shell=True, env=os.environ, cwd=os.getcwd())
124    p.communicate()
125    if p.returncode != 0:
126        sys.exit(p.returncode)
127
128
129def rm(pattern, directory=False):
130    """Recursively remove a file or dir by pattern."""
131    def safe_remove(path):
132        try:
133            os.remove(path)
134        except OSError as err:
135            if err.errno != errno.ENOENT:
136                raise
137        else:
138            safe_print("rm %s" % path)
139
140    def safe_rmtree(path):
141        def onerror(fun, path, excinfo):
142            exc = excinfo[1]
143            if exc.errno != errno.ENOENT:
144                raise
145
146        existed = os.path.isdir(path)
147        shutil.rmtree(path, onerror=onerror)
148        if existed:
149            safe_print("rmdir -f %s" % path)
150
151    if "*" not in pattern:
152        if directory:
153            safe_rmtree(pattern)
154        else:
155            safe_remove(pattern)
156        return
157
158    for root, subdirs, subfiles in os.walk('.'):
159        root = os.path.normpath(root)
160        if root.startswith('.git/'):
161            continue
162        found = fnmatch.filter(subdirs if directory else subfiles, pattern)
163        for name in found:
164            path = os.path.join(root, name)
165            if directory:
166                safe_print("rmdir -f %s" % path)
167                safe_rmtree(path)
168            else:
169                safe_print("rm %s" % path)
170                safe_remove(path)
171
172
173def safe_remove(path):
174    try:
175        os.remove(path)
176    except OSError as err:
177        if err.errno != errno.ENOENT:
178            raise
179    else:
180        safe_print("rm %s" % path)
181
182
183def safe_rmtree(path):
184    def onerror(fun, path, excinfo):
185        exc = excinfo[1]
186        if exc.errno != errno.ENOENT:
187            raise
188
189    existed = os.path.isdir(path)
190    shutil.rmtree(path, onerror=onerror)
191    if existed:
192        safe_print("rmdir -f %s" % path)
193
194
195def recursive_rm(*patterns):
196    """Recursively remove a file or matching a list of patterns."""
197    for root, subdirs, subfiles in os.walk(u'.'):
198        root = os.path.normpath(root)
199        if root.startswith('.git/'):
200            continue
201        for file in subfiles:
202            for pattern in patterns:
203                if fnmatch.fnmatch(file, pattern):
204                    safe_remove(os.path.join(root, file))
205        for dir in subdirs:
206            for pattern in patterns:
207                if fnmatch.fnmatch(dir, pattern):
208                    safe_rmtree(os.path.join(root, dir))
209
210
211# ===================================================================
212# commands
213# ===================================================================
214
215
216def build():
217    """Build / compile"""
218    # Make sure setuptools is installed (needed for 'develop' /
219    # edit mode).
220    sh('%s -c "import setuptools"' % PYTHON)
221
222    # "build_ext -i" copies compiled *.pyd files in ./psutil directory in
223    # order to allow "import psutil" when using the interactive interpreter
224    # from within psutil root directory.
225    cmd = [PYTHON, "setup.py", "build_ext", "-i"]
226    if sys.version_info[:2] >= (3, 6) and os.cpu_count() or 1 > 1:
227        cmd += ['--parallel', str(os.cpu_count())]
228    # Print coloured warnings in real time.
229    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
230    try:
231        for line in iter(p.stdout.readline, b''):
232            if PY3:
233                line = line.decode()
234            line = line.strip()
235            if 'warning' in line:
236                win_colorprint(line, YELLOW)
237            elif 'error' in line:
238                win_colorprint(line, RED)
239            else:
240                print(line)
241        # retcode = p.poll()
242        p.communicate()
243        if p.returncode:
244            win_colorprint("failure", RED)
245            sys.exit(p.returncode)
246    finally:
247        p.terminate()
248        p.wait()
249
250    # Make sure it actually worked.
251    sh('%s -c "import psutil"' % PYTHON)
252    win_colorprint("build + import successful", GREEN)
253
254
255def wheel():
256    """Create wheel file."""
257    build()
258    sh("%s setup.py bdist_wheel" % PYTHON)
259
260
261def upload_wheels():
262    """Upload wheel files on PyPI."""
263    build()
264    sh("%s -m twine upload dist/*.whl" % PYTHON)
265
266
267def install_pip():
268    """Install pip"""
269    try:
270        sh('%s -c "import pip"' % PYTHON)
271    except SystemExit:
272        if PY3:
273            from urllib.request import urlopen
274        else:
275            from urllib2 import urlopen
276
277        if hasattr(ssl, '_create_unverified_context'):
278            ctx = ssl._create_unverified_context()
279        else:
280            ctx = None
281        kw = dict(context=ctx) if ctx else {}
282        safe_print("downloading %s" % GET_PIP_URL)
283        req = urlopen(GET_PIP_URL, **kw)
284        data = req.read()
285
286        tfile = os.path.join(tempfile.gettempdir(), 'get-pip.py')
287        with open(tfile, 'wb') as f:
288            f.write(data)
289
290        try:
291            sh('%s %s --user' % (PYTHON, tfile))
292        finally:
293            os.remove(tfile)
294
295
296def install():
297    """Install in develop / edit mode"""
298    build()
299    sh("%s setup.py develop" % PYTHON)
300
301
302def uninstall():
303    """Uninstall psutil"""
304    # Uninstalling psutil on Windows seems to be tricky.
305    # On "import psutil" tests may import a psutil version living in
306    # C:\PythonXY\Lib\site-packages which is not what we want, so
307    # we try both "pip uninstall psutil" and manually remove stuff
308    # from site-packages.
309    clean()
310    install_pip()
311    here = os.getcwd()
312    try:
313        os.chdir('C:\\')
314        while True:
315            try:
316                import psutil  # NOQA
317            except ImportError:
318                break
319            else:
320                sh("%s -m pip uninstall -y psutil" % PYTHON)
321    finally:
322        os.chdir(here)
323
324    for dir in site.getsitepackages():
325        for name in os.listdir(dir):
326            if name.startswith('psutil'):
327                rm(os.path.join(dir, name))
328            elif name == 'easy-install.pth':
329                # easy_install can add a line (installation path) into
330                # easy-install.pth; that line alters sys.path.
331                path = os.path.join(dir, name)
332                with open(path, 'rt') as f:
333                    lines = f.readlines()
334                    hasit = False
335                    for line in lines:
336                        if 'psutil' in line:
337                            hasit = True
338                            break
339                if hasit:
340                    with open(path, 'wt') as f:
341                        for line in lines:
342                            if 'psutil' not in line:
343                                f.write(line)
344                            else:
345                                print("removed line %r from %r" % (line, path))
346
347
348def clean():
349    """Deletes dev files"""
350    recursive_rm(
351        "$testfn*",
352        "*.bak",
353        "*.core",
354        "*.egg-info",
355        "*.orig",
356        "*.pyc",
357        "*.pyd",
358        "*.pyo",
359        "*.rej",
360        "*.so",
361        "*.~",
362        "*__pycache__",
363        ".coverage",
364        ".failed-tests.txt",
365    )
366    safe_rmtree("build")
367    safe_rmtree(".coverage")
368    safe_rmtree("dist")
369    safe_rmtree("docs/_build")
370    safe_rmtree("htmlcov")
371    safe_rmtree("tmp")
372
373
374def setup_dev_env():
375    """Install useful deps"""
376    install_pip()
377    install_git_hooks()
378    sh("%s -m pip install -U %s" % (PYTHON, " ".join(DEPS)))
379
380
381def lint():
382    """Run flake8 against all py files"""
383    py_files = subprocess.check_output("git ls-files")
384    if PY3:
385        py_files = py_files.decode()
386    py_files = [x for x in py_files.split() if x.endswith('.py')]
387    py_files = ' '.join(py_files)
388    sh("%s -m flake8 %s" % (PYTHON, py_files), nolog=True)
389
390
391def test(name=RUNNER_PY):
392    """Run tests"""
393    build()
394    sh("%s %s" % (PYTHON, name))
395
396
397def coverage():
398    """Run coverage tests."""
399    # Note: coverage options are controlled by .coveragerc file
400    build()
401    sh("%s -m coverage run %s" % (PYTHON, RUNNER_PY))
402    sh("%s -m coverage report" % PYTHON)
403    sh("%s -m coverage html" % PYTHON)
404    sh("%s -m webbrowser -t htmlcov/index.html" % PYTHON)
405
406
407def test_process():
408    """Run process tests"""
409    build()
410    sh("%s psutil\\tests\\test_process.py" % PYTHON)
411
412
413def test_system():
414    """Run system tests"""
415    build()
416    sh("%s psutil\\tests\\test_system.py" % PYTHON)
417
418
419def test_platform():
420    """Run windows only tests"""
421    build()
422    sh("%s psutil\\tests\\test_windows.py" % PYTHON)
423
424
425def test_misc():
426    """Run misc tests"""
427    build()
428    sh("%s psutil\\tests\\test_misc.py" % PYTHON)
429
430
431def test_unicode():
432    """Run unicode tests"""
433    build()
434    sh("%s psutil\\tests\\test_unicode.py" % PYTHON)
435
436
437def test_connections():
438    """Run connections tests"""
439    build()
440    sh("%s psutil\\tests\\test_connections.py" % PYTHON)
441
442
443def test_contracts():
444    """Run contracts tests"""
445    build()
446    sh("%s psutil\\tests\\test_contracts.py" % PYTHON)
447
448
449def test_testutils():
450    """Run test utilities tests"""
451    build()
452    sh("%s psutil\\tests\\test_testutils.py" % PYTHON)
453
454
455def test_by_name(name):
456    """Run test by name"""
457    build()
458    sh("%s -m unittest -v %s" % (PYTHON, name))
459
460
461def test_failed():
462    """Re-run tests which failed on last run."""
463    build()
464    sh("%s %s --last-failed" % (PYTHON, RUNNER_PY))
465
466
467def test_memleaks():
468    """Run memory leaks tests"""
469    build()
470    sh("%s psutil\\tests\\test_memleaks.py" % PYTHON)
471
472
473def install_git_hooks():
474    """Install GIT pre-commit hook."""
475    if os.path.isdir('.git'):
476        src = os.path.join(
477            ROOT_DIR, "scripts", "internal", "git_pre_commit.py")
478        dst = os.path.realpath(
479            os.path.join(ROOT_DIR, ".git", "hooks", "pre-commit"))
480        with open(src, "rt") as s:
481            with open(dst, "wt") as d:
482                d.write(s.read())
483
484
485def bench_oneshot():
486    """Benchmarks for oneshot() ctx manager (see #799)."""
487    sh("%s -Wa scripts\\internal\\bench_oneshot.py" % PYTHON)
488
489
490def bench_oneshot_2():
491    """Same as above but using perf module (supposed to be more precise)."""
492    sh("%s -Wa scripts\\internal\\bench_oneshot_2.py" % PYTHON)
493
494
495def print_access_denied():
496    """Print AD exceptions raised by all Process methods."""
497    build()
498    sh("%s -Wa scripts\\internal\\print_access_denied.py" % PYTHON)
499
500
501def print_api_speed():
502    """Benchmark all API calls."""
503    build()
504    sh("%s -Wa scripts\\internal\\print_api_speed.py" % PYTHON)
505
506
507def download_appveyor_wheels():
508    """Download appveyor wheels."""
509    sh("%s -Wa scripts\\internal\\download_wheels_appveyor.py "
510       "--user giampaolo --project psutil" % PYTHON)
511
512
513def get_python(path):
514    if not path:
515        return sys.executable
516    if os.path.isabs(path):
517        return path
518    # try to look for a python installation given a shortcut name
519    path = path.replace('.', '')
520    vers = (
521        '26',
522        '26-32',
523        '26-64',
524        '27',
525        '27-32',
526        '27-64',
527        '36',
528        '36-32',
529        '36-64',
530        '37',
531        '37-32',
532        '37-64',
533        '38',
534        '38-32',
535        '38-64',
536        '39-32',
537        '39-64',
538    )
539    for v in vers:
540        pypath = r'C:\\python%s\python.exe' % v
541        if path in pypath and os.path.isfile(pypath):
542            return pypath
543
544
545def main():
546    global PYTHON
547    parser = argparse.ArgumentParser()
548    # option shared by all commands
549    parser.add_argument(
550        '-p', '--python',
551        help="use python executable path")
552    sp = parser.add_subparsers(dest='command', title='targets')
553    sp.add_parser('bench-oneshot', help="benchmarks for oneshot()")
554    sp.add_parser('bench-oneshot_2', help="benchmarks for oneshot() (perf)")
555    sp.add_parser('build', help="build")
556    sp.add_parser('clean', help="deletes dev files")
557    sp.add_parser('coverage', help="run coverage tests.")
558    sp.add_parser('download-appveyor-wheels', help="download wheels.")
559    sp.add_parser('help', help="print this help")
560    sp.add_parser('install', help="build + install in develop/edit mode")
561    sp.add_parser('install-git-hooks', help="install GIT pre-commit hook")
562    sp.add_parser('install-pip', help="install pip")
563    sp.add_parser('lint', help="run flake8 against all py files")
564    sp.add_parser('print-access-denied', help="print AD exceptions")
565    sp.add_parser('print-api-speed', help="benchmark all API calls")
566    sp.add_parser('setup-dev-env', help="install deps")
567    test = sp.add_parser('test', help="[ARG] run tests")
568    test_by_name = sp.add_parser('test-by-name', help="<ARG> run test by name")
569    sp.add_parser('test-connections', help="run connections tests")
570    sp.add_parser('test-contracts', help="run contracts tests")
571    sp.add_parser('test-failed', help="re-run tests which failed on last run")
572    sp.add_parser('test-memleaks', help="run memory leaks tests")
573    sp.add_parser('test-misc', help="run misc tests")
574    sp.add_parser('test-platform', help="run windows only tests")
575    sp.add_parser('test-process', help="run process tests")
576    sp.add_parser('test-system', help="run system tests")
577    sp.add_parser('test-unicode', help="run unicode tests")
578    sp.add_parser('test-testutils', help="run test utils tests")
579    sp.add_parser('uninstall', help="uninstall psutil")
580    sp.add_parser('upload-wheels', help="upload wheel files on PyPI")
581    sp.add_parser('wheel', help="create wheel file")
582
583    for p in (test, test_by_name):
584        p.add_argument('arg', type=str, nargs='?', default="", help="arg")
585    args = parser.parse_args()
586
587    # set python exe
588    PYTHON = get_python(args.python)
589    if not PYTHON:
590        return sys.exit(
591            "can't find any python installation matching %r" % args.python)
592    os.putenv('PYTHON', PYTHON)
593    win_colorprint("using " + PYTHON)
594
595    if not args.command or args.command == 'help':
596        parser.print_help(sys.stderr)
597        sys.exit(1)
598
599    fname = args.command.replace('-', '_')
600    fun = getattr(sys.modules[__name__], fname)  # err if fun not defined
601    funargs = []
602    # mandatory args
603    if args.command in ('test-by-name', 'test-script'):
604        if not args.arg:
605            sys.exit('command needs an argument')
606        funargs = [args.arg]
607    # optional args
608    if args.command == 'test' and args.arg:
609        funargs = [args.arg]
610    fun(*funargs)
611
612
613if __name__ == '__main__':
614    main()
615