diff --git a/tools/pip_install.py b/tools/pip_install.py index 7c73a51bb..71584efe5 100755 --- a/tools/pip_install.py +++ b/tools/pip_install.py @@ -11,6 +11,7 @@ from __future__ import absolute_import from __future__ import print_function +import contextlib import os import re import shutil @@ -23,6 +24,17 @@ import readlink import strip_hashes +# Once this code doesn't need to support Python 2, we can simply use +# tempfile.TemporaryDirectory. +@contextlib.contextmanager +def temporary_directory(): + dirpath = tempfile.mkdtemp() + try: + yield dirpath + finally: + shutil.rmtree(dirpath) + + def find_tools_path(): return os.path.dirname(readlink.main(__file__)) @@ -87,12 +99,7 @@ def main(args): tools_path = find_tools_path() working_dir = tempfile.mkdtemp() - if os.environ.get('TRAVIS'): - # When this script is executed on Travis, the following print will make the log - # be folded until the end command is printed (see finally section). - print('travis_fold:start:install_certbot_deps') - - try: + with temporary_directory() as working_dir: test_constraints = os.path.join(working_dir, 'test_constraints.txt') all_constraints = os.path.join(working_dir, 'all_constraints.txt') @@ -121,10 +128,6 @@ def main(args): pip_install_with_print('--constraint "{0}" {1}'.format( all_constraints, ' '.join(args))) - finally: - if os.environ.get('TRAVIS'): - print('travis_fold:end:install_certbot_deps') - shutil.rmtree(working_dir) if __name__ == '__main__': diff --git a/tools/pipstrap.py b/tools/pipstrap.py index 448706ae9..d57d9610d 100755 --- a/tools/pipstrap.py +++ b/tools/pipstrap.py @@ -8,6 +8,7 @@ pinned the same way as our other packages. from __future__ import absolute_import import os +import shutil import tempfile import pip_install @@ -30,7 +31,7 @@ wheel==0.35.1 \ def main(): - with tempfile.TemporaryDirectory() as tempdir: + with pip_install.temporary_directory() as tempdir: requirements_filepath = os.path.join(tempdir, 'reqs.txt') with open(requirements_filepath, 'w') as f: f.write(REQUIREMENTS)