From 67f5eb2c6c46cfabcaafb2491db9fc395e0e69ee Mon Sep 17 00:00:00 2001 From: "Jason R. Coombs" Date: Tue, 3 Feb 2015 22:25:17 +0400 Subject: Use contextlib to simplify run-tests. --- ptr.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/ptr.py b/ptr.py index eefb82f..5ff2995 100644 --- a/ptr.py +++ b/ptr.py @@ -4,9 +4,22 @@ Implementation import os as _os import shlex as _shlex +import contextlib as _contextlib +import sys as _sys import setuptools.command.test as orig +@_contextlib.contextmanager +def _save_argv(repl=None): + saved = _sys.argv[:] + if repl is not None: + _sys.argv[:] = repl + try: + yield saved + finally: + _sys.argv[:] = saved + + class PyTest(orig.test): user_options = [ ('extras', None, "Install (all) setuptools extras when running tests"), @@ -86,12 +99,7 @@ class PyTest(orig.test): def run_tests(self): """ - Override run_tests to invoke pytest. + Invoke pytest, replacing argv. """ - import pytest - import sys - # hide command-line arguments from pytest.main - argv_saved = list(sys.argv) - sys.argv[1:] = self.addopts - self.result_code = pytest.main() - sys.argv[:] = argv_saved + with _save_argv(_sys.argv[:1] + self.addopts): + self.result_code = __import__('pytest').main() -- cgit v1.2.1