1#!/usr/bin/env python
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 errno
16import fnmatch
17import functools
18import os
19import shutil
20import site
21import ssl
22import subprocess
23import sys
24import tempfile
25
26
27PYTHON = os.getenv('PYTHON', sys.executable)
28TSCRIPT = os.getenv('TSCRIPT', 'psutil\\tests\\__main__.py')
29GET_PIP_URL = "https://bootstrap.pypa.io/get-pip.py"
30PY3 = sys.version_info[0] == 3
31DEPS = [
32    "coverage",
33    "flake8",
34    "ipaddress",
35    "mock",
36    "nose",
37    "pdbpp",
38    "perf",
39    "pip",
40    "pypiwin32",
41    "pyreadline",
42    "setuptools",
43    "unittest2",
44    "wheel",
45    "wmi",
46    "requests"
47]
48_cmds = {}
49if PY3:
50    basestring = str
51
52# ===================================================================
53# utils
54# ===================================================================
55
56
57def safe_print(text, file=sys.stdout, flush=False):
58    """Prints a (unicode) string to the console, encoded depending on
59    the stdout/file encoding (eg. cp437 on Windows). This is to avoid
60    encoding errors in case of funky path names.
61    Works with Python 2 and 3.
62    """
63    if not isinstance(text, basestring):
64        return print(text, file=file)
65    try:
66        file.write(text)
67    except UnicodeEncodeError:
68        bytes_string = text.encode(file.encoding, 'backslashreplace')
69        if hasattr(file, 'buffer'):
70            file.buffer.write(bytes_string)
71        else:
72            text = bytes_string.decode(file.encoding, 'strict')
73            file.write(text)
74    file.write("\n")
75
76
77def sh(cmd, nolog=False):
78    if not nolog:
79        safe_print("cmd: " + cmd)
80    p = subprocess.Popen(cmd, shell=True, env=os.environ, cwd=os.getcwd())
81    p.communicate()
82    if p.returncode != 0:
83        sys.exit(p.returncode)
84
85
86def cmd(fun):
87    @functools.wraps(fun)
88    def wrapper(*args, **kwds):
89        return fun(*args, **kwds)
90
91    _cmds[fun.__name__] = fun.__doc__
92    return wrapper
93
94
95def rm(pattern, directory=False):
96    """Recursively remove a file or dir by pattern."""
97    def safe_remove(path):
98        try:
99            os.remove(path)
100        except OSError as err:
101            if err.errno != errno.ENOENT:
102                raise
103        else:
104            safe_print("rm %s" % path)
105
106    def safe_rmtree(path):
107        def onerror(fun, path, excinfo):
108            exc = excinfo[1]
109            if exc.errno != errno.ENOENT:
110                raise
111
112        existed = os.path.isdir(path)
113        shutil.rmtree(path, onerror=onerror)
114        if existed:
115            safe_print("rmdir -f %s" % path)
116
117    if "*" not in pattern:
118        if directory:
119            safe_rmtree(pattern)
120        else:
121            safe_remove(pattern)
122        return
123
124    for root, subdirs, subfiles in os.walk('.'):
125        root = os.path.normpath(root)
126        if root.startswith('.git/'):
127            continue
128        found = fnmatch.filter(subdirs if directory else subfiles, pattern)
129        for name in found:
130            path = os.path.join(root, name)
131            if directory:
132                safe_print("rmdir -f %s" % path)
133                safe_rmtree(path)
134            else:
135                safe_print("rm %s" % path)
136                safe_remove(path)
137
138
139def safe_remove(path):
140    try:
141        os.remove(path)
142    except OSError as err:
143        if err.errno != errno.ENOENT:
144            raise
145    else:
146        safe_print("rm %s" % path)
147
148
149def safe_rmtree(path):
150    def onerror(fun, path, excinfo):
151        exc = excinfo[1]
152        if exc.errno != errno.ENOENT:
153            raise
154
155    existed = os.path.isdir(path)
156    shutil.rmtree(path, onerror=onerror)
157    if existed:
158        safe_print("rmdir -f %s" % path)
159
160
161def recursive_rm(*patterns):
162    """Recursively remove a file or matching a list of patterns."""
163    for root, subdirs, subfiles in os.walk(u'.'):
164        root = os.path.normpath(root)
165        if root.startswith('.git/'):
166            continue
167        for file in subfiles:
168            for pattern in patterns:
169                if fnmatch.fnmatch(file, pattern):
170                    safe_remove(os.path.join(root, file))
171        for dir in subdirs:
172            for pattern in patterns:
173                if fnmatch.fnmatch(dir, pattern):
174                    safe_rmtree(os.path.join(root, dir))
175
176
177def test_setup():
178    os.environ['PYTHONWARNINGS'] = 'all'
179    os.environ['PSUTIL_TESTING'] = '1'
180    os.environ['PSUTIL_DEBUG'] = '1'
181
182
183# ===================================================================
184# commands
185# ===================================================================
186
187
188@cmd
189def help():
190    """Print this help"""
191    safe_print('Run "make [-p <PYTHON>] <target>" where <target> is one of:')
192    for name in sorted(_cmds):
193        safe_print(
194            "    %-20s %s" % (name.replace('_', '-'), _cmds[name] or ''))
195    sys.exit(1)
196
197
198@cmd
199def build():
200    """Build / compile"""
201    # Make sure setuptools is installed (needed for 'develop' /
202    # edit mode).
203    sh('%s -c "import setuptools"' % PYTHON)
204    sh("%s setup.py build" % PYTHON)
205    # Copies compiled *.pyd files in ./psutil directory in order to
206    # allow "import psutil" when using the interactive interpreter
207    # from within this directory.
208    sh("%s setup.py build_ext -i" % PYTHON)
209    # Make sure it actually worked.
210    sh('%s -c "import psutil"' % PYTHON)
211
212
213@cmd
214def build_wheel():
215    """Create wheel file."""
216    build()
217    sh("%s setup.py bdist_wheel" % PYTHON)
218
219
220@cmd
221def install_pip():
222    """Install pip"""
223    try:
224        import pip  # NOQA
225    except ImportError:
226        if PY3:
227            from urllib.request import urlopen
228        else:
229            from urllib2 import urlopen
230
231        if hasattr(ssl, '_create_unverified_context'):
232            ctx = ssl._create_unverified_context()
233        else:
234            ctx = None
235        kw = dict(context=ctx) if ctx else {}
236        safe_print("downloading %s" % GET_PIP_URL)
237        req = urlopen(GET_PIP_URL, **kw)
238        data = req.read()
239
240        tfile = os.path.join(tempfile.gettempdir(), 'get-pip.py')
241        with open(tfile, 'wb') as f:
242            f.write(data)
243
244        try:
245            sh('%s %s --user' % (PYTHON, tfile))
246        finally:
247            os.remove(tfile)
248
249
250@cmd
251def install():
252    """Install in develop / edit mode"""
253    install_git_hooks()
254    build()
255    sh("%s setup.py develop" % PYTHON)
256
257
258@cmd
259def uninstall():
260    """Uninstall psutil"""
261    # Uninstalling psutil on Windows seems to be tricky.
262    # On "import psutil" tests may import a psutil version living in
263    # C:\PythonXY\Lib\site-packages which is not what we want, so
264    # we try both "pip uninstall psutil" and manually remove stuff
265    # from site-packages.
266    clean()
267    install_pip()
268    here = os.getcwd()
269    try:
270        os.chdir('C:\\')
271        while True:
272            try:
273                import psutil  # NOQA
274            except ImportError:
275                break
276            else:
277                sh("%s -m pip uninstall -y psutil" % PYTHON)
278    finally:
279        os.chdir(here)
280
281    for dir in site.getsitepackages():
282        for name in os.listdir(dir):
283            if name.startswith('psutil'):
284                rm(os.path.join(dir, name))
285
286
287@cmd
288def clean():
289    """Deletes dev files"""
290    recursive_rm(
291        "$testfn*",
292        "*.bak",
293        "*.core",
294        "*.egg-info",
295        "*.orig",
296        "*.pyc",
297        "*.pyd",
298        "*.pyo",
299        "*.rej",
300        "*.so",
301        "*.~",
302        "*__pycache__",
303        ".coverage",
304        ".tox",
305    )
306    safe_rmtree("build")
307    safe_rmtree(".coverage")
308    safe_rmtree("dist")
309    safe_rmtree("docs/_build")
310    safe_rmtree("htmlcov")
311    safe_rmtree("tmp")
312
313
314@cmd
315def setup_dev_env():
316    """Install useful deps"""
317    install_pip()
318    install_git_hooks()
319    sh("%s -m pip install -U %s" % (PYTHON, " ".join(DEPS)))
320
321
322@cmd
323def flake8():
324    """Run flake8 against all py files"""
325    py_files = subprocess.check_output("git ls-files")
326    if PY3:
327        py_files = py_files.decode()
328    py_files = [x for x in py_files.split() if x.endswith('.py')]
329    py_files = ' '.join(py_files)
330    sh("%s -m flake8 %s" % (PYTHON, py_files), nolog=True)
331
332
333@cmd
334def test():
335    """Run tests"""
336    install()
337    test_setup()
338    sh("%s %s" % (PYTHON, TSCRIPT))
339
340
341@cmd
342def coverage():
343    """Run coverage tests."""
344    # Note: coverage options are controlled by .coveragerc file
345    install()
346    test_setup()
347    sh("%s -m coverage run %s" % (PYTHON, TSCRIPT))
348    sh("%s -m coverage report" % PYTHON)
349    sh("%s -m coverage html" % PYTHON)
350    sh("%s -m webbrowser -t htmlcov/index.html" % PYTHON)
351
352
353@cmd
354def test_process():
355    """Run process tests"""
356    install()
357    test_setup()
358    sh("%s -m unittest -v psutil.tests.test_process" % PYTHON)
359
360
361@cmd
362def test_system():
363    """Run system tests"""
364    install()
365    test_setup()
366    sh("%s -m unittest -v psutil.tests.test_system" % PYTHON)
367
368
369@cmd
370def test_platform():
371    """Run windows only tests"""
372    install()
373    test_setup()
374    sh("%s -m unittest -v psutil.tests.test_windows" % PYTHON)
375
376
377@cmd
378def test_misc():
379    """Run misc tests"""
380    install()
381    test_setup()
382    sh("%s -m unittest -v psutil.tests.test_misc" % PYTHON)
383
384
385@cmd
386def test_unicode():
387    """Run unicode tests"""
388    install()
389    test_setup()
390    sh("%s -m unittest -v psutil.tests.test_unicode" % PYTHON)
391
392
393@cmd
394def test_connections():
395    """Run connections tests"""
396    install()
397    test_setup()
398    sh("%s -m unittest -v psutil.tests.test_connections" % PYTHON)
399
400
401@cmd
402def test_contracts():
403    """Run contracts tests"""
404    install()
405    test_setup()
406    sh("%s -m unittest -v psutil.tests.test_contracts" % PYTHON)
407
408
409@cmd
410def test_by_name():
411    """Run test by name"""
412    try:
413        safe_print(sys.argv)
414        name = sys.argv[2]
415    except IndexError:
416        sys.exit('second arg missing')
417    install()
418    test_setup()
419    sh("%s -m unittest -v %s" % (PYTHON, name))
420
421
422@cmd
423def test_script():
424    """Quick way to test a script"""
425    try:
426        safe_print(sys.argv)
427        name = sys.argv[2]
428    except IndexError:
429        sys.exit('second arg missing')
430    install()
431    test_setup()
432    sh("%s %s" % (PYTHON, name))
433
434
435@cmd
436def test_memleaks():
437    """Run memory leaks tests"""
438    install()
439    test_setup()
440    sh("%s psutil\\tests\\test_memory_leaks.py" % PYTHON)
441
442
443@cmd
444def install_git_hooks():
445    if os.path.isdir('.git'):
446        shutil.copy(".git-pre-commit", ".git\\hooks\\pre-commit")
447
448
449@cmd
450def bench_oneshot():
451    install()
452    sh("%s -Wa scripts\\internal\\bench_oneshot.py" % PYTHON)
453
454
455@cmd
456def bench_oneshot_2():
457    install()
458    sh("%s -Wa scripts\\internal\\bench_oneshot_2.py" % PYTHON)
459
460
461def set_python(s):
462    global PYTHON
463    if os.path.isabs(s):
464        PYTHON = s
465    else:
466        # try to look for a python installation
467        orig = s
468        s = s.replace('.', '')
469        vers = ('26', '27', '34', '35', '36', '37',
470                '26-64', '27-64', '34-64', '35-64', '36-64', '37-64')
471        for v in vers:
472            if s == v:
473                path = 'C:\\python%s\python.exe' % s
474                if os.path.isfile(path):
475                    print(path)
476                    PYTHON = path
477                    os.putenv('PYTHON', path)
478                    return
479        return sys.exit(
480            "can't find any python installation matching %r" % orig)
481
482
483def parse_cmdline():
484    if '-p' in sys.argv:
485        try:
486            pos = sys.argv.index('-p')
487            sys.argv.pop(pos)
488            py = sys.argv.pop(pos)
489        except IndexError:
490            return help()
491        set_python(py)
492
493
494def main():
495    parse_cmdline()
496    try:
497        cmd = sys.argv[1].replace('-', '_')
498    except IndexError:
499        return help()
500    if cmd in _cmds:
501        fun = getattr(sys.modules[__name__], cmd)
502        fun()
503    else:
504        help()
505
506
507if __name__ == '__main__':
508    main()
509