From d105606fe67f565f815029f456ccd1f548b8fe61 Mon Sep 17 00:00:00 2001 From: Eddie Lebow Date: Mon, 6 Apr 2020 21:55:24 -0400 Subject: [PATCH] Break out `utils.run_shell_command` to new class --- augur/util_support/__init__.py | 0 augur/util_support/shell_command_runner.py | 90 +++++++++++++++++++ augur/utils.py | 68 ++------------ .../util_support/test_shell_command_runner.py | 64 +++++++++++++ 4 files changed, 160 insertions(+), 62 deletions(-) create mode 100644 augur/util_support/__init__.py create mode 100644 augur/util_support/shell_command_runner.py create mode 100644 tests/util_support/test_shell_command_runner.py diff --git a/augur/util_support/__init__.py b/augur/util_support/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/augur/util_support/shell_command_runner.py b/augur/util_support/shell_command_runner.py new file mode 100644 index 000000000..40ba953fc --- /dev/null +++ b/augur/util_support/shell_command_runner.py @@ -0,0 +1,90 @@ +import os +import sys +import subprocess +from textwrap import dedent + + +class ShellCommandRunner: + """ + Run the given command string via Bash with error checking. + + TODO move to method docstrings + Returns True if the command exits normally. Returns False if the command + exits with failure and "raise_errors" is False (the default). When + "raise_errors" is True, exceptions are rethrown. + + If an *extra_env* mapping is passed, the provided keys and values are + overlayed onto the default subprocess environment. + """ + + def __init__(self, cmd, *, raise_errors=False, extra_env=None): + self.cmd = cmd + self.raise_errors = raise_errors + self.extra_env = extra_env + + def run(self): + try: + self.invoke_command() + except Exception as error: + self.print_error_message(error) + + if self.raise_errors: + raise error + + return False + + return True + + def invoke_command(self): + return subprocess.check_output( + self.shell_executable + self.shell_args, + shell=False, + stderr=subprocess.STDOUT, + env=self.modified_env, + ) + + @property + def shell_executable(self): + if os.name == "posix": + return ["/bin/bash"] + else: + # We try best effort on other systems. For now that means nt/java. + return ["env", "bash"] + + @property + def shell_args(self): + return ["-c", "set -euo pipefail; " + self.cmd] + + @property + def modified_env(self): + env = os.environ.copy() + + if self.extra_env: + env.update(self.extra_env) + + return env + + def print_error_message(self, error): + if isinstance(error, subprocess.CalledProcessError): + message = f"{error.output}\nshell exited {error.returncode} when running: {self.cmd}" + + if error.returncode == 127: + message += "\nAre you sure this program is installed?" + elif isinstance(error, FileNotFoundError): + shell = " and ".join(self.shell_executable) + + message = f""" + Unable to run shell commands using {shell}! + + Augur requires {shell} to be installed. Please open an issue on GitHub + if you need assistance. + """ + else: + message = str(error) + + self.print_error(message) + + @staticmethod + def print_error(message): + """Prints message to STDERR formatted with textwrap.dedent""" + print("\nERROR: " + dedent(message).lstrip("\n") + "\n", file=sys.stderr) diff --git a/augur/utils.py b/augur/utils.py index ef7d97f36..a63aeaf3b 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -11,14 +11,17 @@ from collections import defaultdict from pkg_resources import resource_stream from io import TextIOWrapper -from textwrap import dedent from .__version__ import __version__ import packaging.version as packaging_version from .validate import validate, ValidateError, load_json_schema +from augur.util_support.shell_command_runner import ShellCommandRunner + + class AugurException(Exception): pass + @contextmanager def open_file(fname, mode): """Open a file using either gzip.open() or open() depending on file name. Semantics identical to open()""" @@ -561,7 +564,7 @@ def write_VCF_translation(prot_dict, vcf_file_name, ref_file_name): shquote = shlex.quote -def run_shell_command(cmd, raise_errors = False, extra_env = None): +def run_shell_command(cmd, raise_errors=False, extra_env=None): """ Run the given command string via Bash with error checking. @@ -572,66 +575,7 @@ def run_shell_command(cmd, raise_errors = False, extra_env = None): If an *extra_env* mapping is passed, the provided keys and values are overlayed onto the default subprocess environment. """ - env = os.environ.copy() - - if extra_env: - env.update(extra_env) - - shargs = ['-c', "set -euo pipefail; " + cmd] - - if os.name == 'posix': - shellexec = ['/bin/bash'] - else: - # We try best effort on other systems. For now that means nt/java. - shellexec = ['env', 'bash'] - - try: - # Use check_call() instead of run() since the latter was added only in Python 3.5. - subprocess.check_output( - shellexec + shargs, - shell = False, - stderr = subprocess.STDOUT, - env = env) - - except subprocess.CalledProcessError as error: - print_error( - "{out}\nshell exited {rc} when running: {cmd}{extra}", - out = error.output, - rc = error.returncode, - cmd = cmd, - extra = "\nAre you sure this program is installed?" if error.returncode==127 else "", - ) - if raise_errors: - raise - else: - return False - - except FileNotFoundError as error: - print_error( - """ - Unable to run shell commands using {shell}! - - Augur requires {shell} to be installed. Please open an issue on GitHub - if you need assistance. - """, - shell = ' and '.join(shellexec) - ) - if raise_errors: - raise - else: - return False - - else: - return True - - -def print_error(message, **kwargs): - """ - Formats *message* with *kwargs* using :meth:`str.format` and - :func:`textwrap.dedent` and uses it to print an error message to - ``sys.stderr``. - """ - print("\nERROR: " + dedent(message.format(**kwargs)).lstrip("\n")+"\n", file = sys.stderr) + return ShellCommandRunner(cmd, raise_errors=raise_errors, extra_env=extra_env).run() def first_line(text): diff --git a/tests/util_support/test_shell_command_runner.py b/tests/util_support/test_shell_command_runner.py new file mode 100644 index 000000000..50b9bd6db --- /dev/null +++ b/tests/util_support/test_shell_command_runner.py @@ -0,0 +1,64 @@ +import os +import re +import subprocess + +from augur.util_support.shell_command_runner import ShellCommandRunner + +import pytest + + +class TestShellCommandRunner: + def test_run_exception_no_raise(self, mocker): + mocker.patch( + "subprocess.check_output", + side_effect=subprocess.CalledProcessError(5, "actual-cmd"), + ) + + assert ShellCommandRunner("great-command", raise_errors=False).run() is False + + def test_run_exception_raise(self, mocker): + mocker.patch( + "subprocess.check_output", + side_effect=subprocess.CalledProcessError(5, "actual-cmd"), + ) + + with pytest.raises(subprocess.CalledProcessError): + ShellCommandRunner("great-command", raise_errors=True).run() + + def test_modified_env(self): + modified_env = ShellCommandRunner( + "great-command", extra_env={"a": 5} + ).modified_env + + assert modified_env["a"] == 5 + assert modified_env["HOME"] == os.getenv("HOME") + + @pytest.mark.parametrize( + "exception, expected_message", + [ + ( + subprocess.CalledProcessError(5, "actual-cmd", output="some error"), + "some error.*shell exited 5 when running: cmd", + ), + ( + FileNotFoundError(), + "Unable to run shell commands using /bin/bash", + ), + ( + Exception("generic or other exception"), + "generic or other exception", + ) + ] + ) + def test_print_error_message(self, mocker, exception, expected_message): + mock_print_error = mocker.patch( + "augur.util_support.shell_command_runner.ShellCommandRunner.print_error" + ) + + ShellCommandRunner("cmd").print_error_message(exception) + + assert re.search( + expected_message, + mock_print_error.call_args[0][0], + re.MULTILINE | re.DOTALL + )