wrapperutils.py 14.7 KB
Newer Older
1
2
#!/usr/bin/env python
#
Paul McCarthy's avatar
Paul McCarthy committed
3
4
# wrapperutils.py - Functions and decorators used by the FSL wrapper
# functions.
5
6
7
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
Paul McCarthy's avatar
Paul McCarthy committed
8
9
10
11
12
13
14
15
16
17
18
"""This module contains functions and decorators used by the FSL wrapper
functions.

.. autosummary::
   :nosignatures:

   applyArgStyle
   required
   fileOrImage
   fileOrArray
"""
19
20
21
22


import os.path as op
import            os
Paul McCarthy's avatar
Paul McCarthy committed
23
import            sys
24
25
26
27
28
29
30
31
import            inspect
import            tempfile
import            warnings
import            functools
import            collections

import            six
import nibabel as nib
Paul McCarthy's avatar
Paul McCarthy committed
32
import numpy   as np
33
34
35
36
37

import fsl.utils.tempdir as tempdir
import fsl.data.image    as fslimage


Paul McCarthy's avatar
Paul McCarthy committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def _update_wrapper(wrapper, wrapped, *args, **kwargs):
    """Replacement for the built-in ``functools.update_wrapper``. This
    implementation ensures that the wrapper function has an attribute
    called ``__wrapped__``, which refers to the ``wrapped`` function.

    This behaviour is only required in Python versions < 3.4.
    """

    wrapper = functools.update_wrapper(wrapper, wrapped, *args, **kwargs)

    # Python >= 3.4 does things right
    if sys.version_info[0] * 10 + sys.version_info[1] < 3.4:
        wrapper.__wrapped__ = wrapped
    return wrapper
52
53


Paul McCarthy's avatar
Paul McCarthy committed
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def _unwrap(func):
    """Replacement for the built-in ``inspect.unwrap`` function, which
    is not present in Python versions prior to 3.4.
    """

    # Python >= 3.4 has an inspect.unwrap function
    if sys.version_info[0] * 10 + sys.version_info[1] < 3.4:
        return inspect.unwrap(func)

    # Otherwise we follow the __wrapped__ chain ourselves
    if hasattr(func, '__wrapped__'):
        return _unwrap(func.__wrapped__)

    return func



SHOW_IF_TRUE = object()
"""Constant to be used in the ``valmap`` passed to the :func:`applyArgStyle`
function.

When a ``SHOW_IF_TRUE`` argument is ``True``, it is added to the generated
command line arguments.
"""


HIDE_IF_TRUE = object()
"""Constant to be used in the ``valmap`` passed to the :func:`applyArgStyle`
function.

When a ``HIDE_IF_TRUE`` argument is ``True``, it is suppressed from the
generated command line arguments.
"""
87
88
89


def applyArgStyle(style, argmap=None, valmap=None, **kwargs):
Paul McCarthy's avatar
Paul McCarthy committed
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    """Turns the given ``kwargs`` into command line options. This function
    is intended to be used to automatically generate command line options
    from arguments passed into a Python function.

    :arg style:  Controls how the ``kwargs`` are converted into command-line
                 options - must be one of the following:
                  - `'-'`: ``-name val``
                  - `'--'`: ``--name val``
                  - `'-='`: ``-name=val``
                  - `'--='`: ``--name=val``

    :arg argmap: Dictionary of ``{kwarg-name : cli-name}`` mappings. This can
                 be used if you want to use different argument names in your
                 Python function for the command-line options.

    :arg valmap: Dictionary of ``{cli-name : value}`` mappings. This can be
                 used to define specific semantics for some command-line
                 options. Acceptable values for ``value`` are as follows

                  - :data:`SHOW_IF_TRUE` - if the argument is present, and
                    ``True`` in ``kwargs``, the command line option
                    will be added (without any arguments).

                  - :data:`HIDE_IF_TRUE` - if the argument is present, and
                    ``False`` in ``kwargs``, the command line option
                    will be added (without any arguments).

                  - Any other constant value. If the argument is present
                    in ``kwargs``, its command-line option will be
                    added, with the constant value as its argument.

                 The argument for any options not specified in the ``valmap``
                 will be converted into strings.
123

Paul McCarthy's avatar
Paul McCarthy committed
124
125
126
127
128
129
130
131
132
133
134
135
    :arg kwargs: Arguments to be converted into command-line options.

    :returns:    A list containing the generated command-line options.
    """

    if style not in ('-', '--', '-=', '--='):
        raise ValueError('Invalid style: {}'.format(style))

    if argmap is None: argmap = {}
    if valmap is None: valmap = {}

    def fmtarg(arg):
136
137
138
139
        if   style in ('-',  '-='):  arg =  '-{}'.format(arg)
        elif style in ('--', '--='): arg = '--{}'.format(arg)
        return arg

Paul McCarthy's avatar
Paul McCarthy committed
140
    def fmtval(val):
141
142
143
144
145
146
147
148
149
150
151
        if     isinstance(val, collections.Sequence) and \
           not isinstance(val, six.string_types):
            return ' '.join([str(v) for v in val])
        else:
            return str(val)

    args = []

    for k, v in kwargs.items():

        k    = argmap.get(k, k)
Paul McCarthy's avatar
Paul McCarthy committed
152
153
154
155
156
157
        mapv = valmap.get(k, fmtval(v))
        k    = fmtarg(k)

        if (mapv is SHOW_IF_TRUE and     v) or \
           (mapv is HIDE_IF_TRUE and not v):
            args.append(k)
158
159
160
161
162
163
164
165
166
167
168
169
170

        elif '=' in style:
            args.append('{}={}'.format(k, mapv))
        else:
            args.extend((k, mapv))

    return args


def required(*reqargs):
    """Decorator which makes sure that all specified keyword arguments are
    present before calling the decorated function.
    """
Paul McCarthy's avatar
Paul McCarthy committed
171

172
    def decorator(func):
Paul McCarthy's avatar
Paul McCarthy committed
173
174
175
        def wrapper(*args, **kwargs):
            kwargs = kwargs.copy()
            kwargs.update(argsToKwargs(func, args))
176
177
178
            for reqarg in reqargs:
                assert reqarg in kwargs
            return func(**kwargs)
Paul McCarthy's avatar
Paul McCarthy committed
179
180
        return _update_wrapper(wrapper, func)

181
182
183
184
185
186
    return decorator


def argsToKwargs(func, args):
    """Given a function, and a sequence of positional arguments destined
    for that function, converts the positional arguments into a dict
Paul McCarthy's avatar
Paul McCarthy committed
187
    of keyword arguments. Used by the :class:`_FileOrThing` class.
188
    """
Paul McCarthy's avatar
Paul McCarthy committed
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204

    func = _unwrap(func)

    # getargspec is the only way to get the names
    # of positional arguments in Python 2.x.
    if sys.version_info[0] < 3:
        argnames = inspect.getargspec(func).args

    # getargspec is deprecated in python 3.x
    else:

        # getfullargspec is deprecated in
        # python 3.5, but not in python 3.6.
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore', category=DeprecationWarning)
            argnames = inspect.getfullargspec(func).args
205
206

    kwargs = collections.OrderedDict()
Paul McCarthy's avatar
Paul McCarthy committed
207
    for name, val in zip(argnames, args):
208
209
210
211
212
213
        kwargs[name] = val

    return kwargs


RETURN = object()
Paul McCarthy's avatar
Paul McCarthy committed
214
215
"""Constant used by the :class:`_FileOrThing` class to indicate that an output
file should be loaded into memory and returned as a Python object.
216
217
218
"""


Paul McCarthy's avatar
Paul McCarthy committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
class _FileOrThing(object):
    """Decorator which ensures that certain arguments which are passed into the
    decorated function are always passed as file names. Both positional and
    keyword arguments can be specified.


    The ``_FileOrThing`` class is not intended to be used directly - see the
    :func:`fileOrImage` and :func:`fileOrArray` decorator functions for more
    details.

    These decorators are intended for functions which wrap a command-line tool,
    i.e. where some inputs/outputs need to be specified as file names.


    **Inputs**


    Any arguments which are not of type ``Thing`` are passed through to the
    decorated function unmodified.  Arguments which are of type ``Thing`` are
    saved to a temporary file, and the name of that file is passed to the
    function.

241

Paul McCarthy's avatar
Paul McCarthy committed
242
    **Outputs**
243
244


Paul McCarthy's avatar
Paul McCarthy committed
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
    If an argument is given the special :data:`RETURN` value, it is assumed
    to be an output argument. In this case, it is replaced with a temporary
    file name then, after the function has completed, that file is loaded
    into memory, and the value returned (along with the function's output,
    and any other arguments with a value of ``RETURN``).


    **Return value**


    Functions decorated with a ``_FileOrThing`` decorator will always return a
    tuple, where the first element is the function's actual return value. The
    remainder of the tuple will contain any arguments that were given the
    special ``RETURN`` value. ``None`` is returned for any ``RETURN``
    arguments corresponded to output files that were not generated by the
    function.


    **Example**


    As an example of using the ``fileOrArray`` decorator on a function
    which concatenates two files containing affine transformations, and
    saves the output to a file::

        # if atob, btoc, or output are passed
        # in as arrays, they are converted to
        # file names.
        @fileOrArray('atob', 'btoc', 'output')
        def concat(atob, btoc, output=None):

            # inputs are guaranteed to be files
            atob = np.loadtxt(atob)
            btoc = np.loadtxt(atoc)

            atoc = np.dot(btoc, atob)

            if output is not None:
                np.savetxt(output, atoc)


    Because we have decorated the ``concat`` function with :func:`fileToArray`,
    it can be called with either file names, or Numpy arrays::

        # All arguments are passed through
        # unmodified - the output will be
        # saved to a file called atoc.mat
        concat('atob.txt', 'btoc.txt', 'atoc.mat')

        # The output is returned as a numpy
        # array (in a tuple with the concat
        # function's return value)
        atoc = concat('atob.txt', 'btoc.txt', RETURN)[1]

        # The inputs are saved to temporary
        # files, and those file names are
        # passed to the concat function.
        atoc = concat(np.diag([2, 2, 2, 0]), np.diag([3, 3, 3, 3]), RETURN)[1]
303
304
305
    """


Paul McCarthy's avatar
Paul McCarthy committed
306
307
308
309
310
311
312
313
    def __init__(self, prepareThing, loadThing, *things):
        """Initialise a ``_FileOrThing`` decorator.

        :arg prepareThing: Function which
        :arg loadThing:    Function which is called for arguments that
                           were set to :data:`RETURN`.

        :arg things:
314
        """
Paul McCarthy's avatar
Paul McCarthy committed
315
316
317
        self.__prepareThing = prepareThing
        self.__loadThing    = loadThing
        self.__things       = things
318
319
320


    def __call__(self, func):
Paul McCarthy's avatar
Paul McCarthy committed
321
322
323
324
325
326
327
328
329
330
331
        """Creates and returns the real decorator function. """

        self.__func          = func
        self.__isFileOrThing = False

        if hasattr(func, '__self__'):
            self.__isFileOrThing = isinstance(func.__self__, _FileOrThing)

        wrapper = functools.partial(self.__wrapper, func)

        return _update_wrapper(wrapper, func)
332
333
334


    def __wrapper(self, func, *args, **kwargs):
Paul McCarthy's avatar
Paul McCarthy committed
335
336
337
        """Function which wraps ``func``, ensuring that any arguments of
        type ``Thing`` are saved to temporary files, and any arguments
        with the value :data:`RETURN` are loaded and returned.
338
339
        """

Paul McCarthy's avatar
Paul McCarthy committed
340
341
342
343
        func          = self.__func
        isFileOrThing = self.__isFileOrThing

        kwargs = kwargs.copy()
344
345
346
        kwargs.update(argsToKwargs(func, args))

        # Create a tempdir to store any temporary
Paul McCarthy's avatar
Paul McCarthy committed
347
        # input/output things, but don't change
348
349
350
351
        # into it, as file paths passed to the
        # function may be relative.
        with tempdir.tempdir(changeto=False) as td:

Paul McCarthy's avatar
Paul McCarthy committed
352
            kwargs, infiles, outfiles = self.__prepareThings(td, kwargs)
353
354

            # Call the function
Paul McCarthy's avatar
Paul McCarthy committed
355
            result = func(**kwargs)
356

Paul McCarthy's avatar
Paul McCarthy committed
357
            # Load the output things that
358
            # were specified as RETURN
Paul McCarthy's avatar
Paul McCarthy committed
359
            outthings = []
360
361
362
363
            for of in outfiles:

                # output file didn't get created
                if not op.exists(of):
Paul McCarthy's avatar
Paul McCarthy committed
364
                    ot = None
365

Paul McCarthy's avatar
Paul McCarthy committed
366
                # load the thing
367
                else:
Paul McCarthy's avatar
Paul McCarthy committed
368
                    ot = self.__loadThing(of)
369

Paul McCarthy's avatar
Paul McCarthy committed
370
                outthings.append(ot)
371

Paul McCarthy's avatar
Paul McCarthy committed
372
373
374
375
376
377
            if isFileOrThing:
                things = result[1:]
                result = result[0]
                return tuple([result] + things + outthings)
            else:
                return tuple([result] + outthings)
378
379


Paul McCarthy's avatar
Paul McCarthy committed
380
    def __prepareThings(self, workdir, kwargs):
381
382
383
384
385
386
387
        """
        """

        kwargs   = dict(kwargs)
        infiles  = []
        outfiles = []

Paul McCarthy's avatar
Paul McCarthy committed
388
        for tname in self.__things:
389

Paul McCarthy's avatar
Paul McCarthy committed
390
            tval = kwargs.get(tname, None)
391

Paul McCarthy's avatar
Paul McCarthy committed
392
            if tval is None:
393
394
                continue

Paul McCarthy's avatar
Paul McCarthy committed
395
396
397
398
399
400
            tval, infile, outfile = self.__prepareThing(workdir, tname, tval)

            if infile  is not None: infiles .append(infile)
            if outfile is not None: outfiles.append(outfile)

            kwargs[tname] = tval
401
402
403
404

        return kwargs, infiles, outfiles


Paul McCarthy's avatar
Paul McCarthy committed
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
def fileOrImage(*imgargs):
    """Decorator which can be used to ensure that any NIfTI images are saved
    to file, and output images can be loaded and returned as ``nibabel``
    image objects.
    """

    def prepareArg(workdir, name, val):

        newval  = val
        infile  = None
        outfile = None

        # This is an input image which has
        # been specified as an in-memory
        # nibabel image. if the image has
        # a backing file, replace the image
        # object with the file name.
        # Otherwise, save the image out to
        # a temporary file, and replace the
        # image with the file name.
        if isinstance(val, nib.nifti1.Nifti1Image):
            imgfile = val.get_filename()

            # in-memory image - we have
            # to save it out to a file
            if imgfile is None:

                hd, imgfile = tempfile.mkstemp(fslimage.defaultExt())

                os.close(hd)
                val.to_filename(imgfile)
                infile = imgfile

            # replace the image with its
            # file name
            newval = imgfile

        # This is an output image, and the
        # caller has requested that it be
        # returned from the function call
        # as an in-memory image.
        elif val == RETURN:
            newval  = op.join(workdir, '{}.nii.gz'.format(name))
            outfile = newval

        return newval, infile, outfile

    def loadImage(path):
        # create an independent in-memory
        # copy of the image file
        img = nib.load(path)
        return nib.nifti1.Nifti1Image(img.get_data(), None, img.header)

    return _FileOrThing(prepareArg, loadImage, *imgargs)


def fileOrArray(*arrargs):
    """Decorator which can be used to ensure that any Numpy arrays are saved
    to text files, and output files can be loaded and returned as Numpy arrays.
    """

    def prepareArg(workdir, name, val):

        newval  = val
        infile  = None
        outfile = None

        # Input has been provided as a numpy
        # array - save it to a file, and
        # replace the argument with the file
        # name
        if isinstance(val, np.ndarray):

            hd, arrfile = tempfile.mkstemp('.txt')

            os.close(hd)

            np.savetxt(arrfile, val, fmt='%0.18f')
            newval = arrfile

        # This is an output, and the caller has
        # requested that it be returned from the
        # function call as an in-memory array.
        elif val == RETURN:
            newval  = op.join(workdir, '{}.txt'.format(name))
            outfile = newval

        return newval, infile, outfile

    return _FileOrThing(prepareArg, np.loadtxt, *arrargs)