run.py 5.22 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
#!/usr/bin/env python
#
# run.py - Functions for running shell commands
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides some functions for running shell commands.

.. autosummary::
   :nosignatures:

   run
   runfsl
14
15
   wait
   dryrun
16
17
18
19
"""


import               logging
20
import               contextlib
21
import               collections
22
23
24
import subprocess as sp
import os.path    as op

25
26
import               six

27
28
from   fsl.utils.platform import platform as fslplatform
import fsl.utils.fslsub                   as fslsub
29
30
31
32
33


log = logging.getLogger(__name__)


34
35
36
37
38
39
DRY_RUN = False
"""If ``True``, the :func:`run` function will only log commands, but will not
execute them.
"""


40
41
42
43
44
45
46
class FSLNotPresent(Exception):
    """Error raised by the :func:`runfsl` function when ``$FSLDIR`` cannot
    be found.
    """
    pass


47
48
49
50
@contextlib.contextmanager
def dryrun(*args):
    """Context manager which causes all calls to :func:`run` to be logged but
    not executed. See the :data:`DRY_RUN` flag.
51
52

    The returned standard output will be equal to ``' '.join(args)``.
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    """
    global DRY_RUN

    oldval  = DRY_RUN
    DRY_RUN = True

    try:
        yield
    finally:
        DRY_RUN = oldval


def _prepareArgs(args):
    """Used by the :func:`run` function. Ensures that the given arguments is a
    list of strings.
68
69
70
71
    """

    if len(args) == 1:

72
73
74
        # Argument was a command string
        if isinstance(args[0], six.string_types):
            args = args[0].split()
75

76
77
78
        # Argument was an unpacked sequence
        else:
            args = args[0]
79

80
81
82
    return list(args)


83
def run(*args, **kwargs):
84
85
    """Call a command and return its output. You can pass the command and
    arguments as a single string, or as a regular or unpacked sequence.
86

87
88
89
    The command can be run on a cluster by using the ``submit`` keyword
    argument.

90
    An exception is raised if the command returns a non-zero exit code, unless
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    the ``ret`` option is set to ``True``.

    :arg submit: Must be passed as a keyword argument. Defaults to ``None``.
                 Accepted values are ``True`` or a
                 If ``True``, the command is submitted as a cluster job via
                 the :func:`.fslsub.submit` function.  May also be a
                 dictionary containing arguments to that function.

    :arg err:    Must be passed as a keyword argument. Defaults to
                 ``False``. If ``True``, standard error is captured and
                 returned. Ignored if ``submit`` is specified.

    :arg ret:    Must be passed as a keyword argument. Defaults to ``False``.
                 If ``True``, and the command's return code is non-0, an
                 exception is not raised.  Ignored if ``submit`` is specified.

    :returns:    If ``submit`` is provided, the cluster job ID is returned.
                 Otherwise if ``err is False and ret is False`` (the default)
                 a string containing the command's standard output.  is
                 returned. Or, if ``err is True`` and/or ``ret is True``, a
                 tuple containing the standard output, standard error (if
                 ``err``), and return code (if ``ret``).
    """
114

115
116
117
    err    = kwargs.get('err',    False)
    ret    = kwargs.get('ret',    False)
    submit = kwargs.get('submit', None)
118
    args   = _prepareArgs(args)
119

120
121
    if not bool(submit):
        submit = None
122

123
124
125
126
127
128
    if submit is not None:
        err = False
        ret = False

        if submit is True:
            submit = dict()
129

130
131
132
    if submit is not None and not isinstance(submit, collections.Mapping):
        raise ValueError('submit must be a mapping containing '
                         'options for fsl.utils.fslsub.submit')
133
134
135
136
137
138
139

    if DRY_RUN:
        log.debug('dryrun: {}'.format(' '.join(args)))
    else:
        log.debug('run: {}'.format(' '.join(args)))

    if DRY_RUN:
140
        stderr = ''
Paul McCarthy's avatar
Paul McCarthy committed
141
        if submit is None:
142
143
144
145
146
147
148
            stdout = ' '.join(args)
        else:
            stdout = '[submit] ' + ' '.join(args)

    elif submit is not None:
        return fslsub.submit(' '.join(args), **submit)

149
    else:
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        proc           = sp.Popen(args, stdout=sp.PIPE, stderr=sp.PIPE)
        stdout, stderr = proc.communicate()
        retcode        = proc.returncode

        stdout = stdout.decode('utf-8').strip()
        stderr = stderr.decode('utf-8').strip()

        log.debug('stdout: {}'.format(stdout))
        log.debug('stderr: {}'.format(stderr))

        if not ret and (retcode != 0):
            raise RuntimeError('{} returned non-zero exit code: {}'.format(
                args[0], retcode))

    results = [stdout]
165

166
167
    if err: results.append(stderr)
    if ret: results.append(retcode)
168

169
170
    if len(results) == 1: return results[0]
    else:                 return tuple(results)
171
172


173
def runfsl(*args, **kwargs):
174
    """Call a FSL command and return its output. This function simply prepends
175
    ``$FSLDIR/bin/`` to the command before passing it to :func:`run`.
176
177
178
    """

    if fslplatform.fsldir is None:
179
        raise FSLNotPresent('$FSLDIR is not set - FSL cannot be found!')
180

181
    args    = _prepareArgs(args)
182
183
    args[0] = op.join(fslplatform.fsldir, 'bin', args[0])

184
    return run(*args, **kwargs)
185
186


187
def wait(job_ids):
188
    """Proxy for :func:`.fslsub.wait`. """
189
    return fslsub.wait(job_ids)