wrapperutils.py 15.2 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)
179
180
181
182
183
184
185
186
187
188
189

        wrapper = _update_wrapper(wrapper, func)

        # If this is a bound method, make
        # sure that the instance is set on
        # the wrapper function - this is
        # needed by _FileOrThing decorators.
        if hasattr(func, '__self__'):
            wrapper.__self__ = func.__self__

        return wrapper
Paul McCarthy's avatar
Paul McCarthy committed
190

191
192
193
194
195
196
    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
197
    of keyword arguments. Used by the :class:`_FileOrThing` class.
198
    """
Paul McCarthy's avatar
Paul McCarthy committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214

    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
215
216

    kwargs = collections.OrderedDict()
Paul McCarthy's avatar
Paul McCarthy committed
217
    for name, val in zip(argnames, args):
218
219
220
221
222
        kwargs[name] = val

    return kwargs


Paul McCarthy's avatar
Paul McCarthy committed
223
LOAD = object()
Paul McCarthy's avatar
Paul McCarthy committed
224
225
"""Constant used by the :class:`_FileOrThing` class to indicate that an output
file should be loaded into memory and returned as a Python object.
226
227
228
"""


Paul McCarthy's avatar
Paul McCarthy committed
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
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.

251

Paul McCarthy's avatar
Paul McCarthy committed
252
    **Outputs**
253
254


Paul McCarthy's avatar
Paul McCarthy committed
255
    If an argument is given the special :data:`LOAD` value, it is assumed
Paul McCarthy's avatar
Paul McCarthy committed
256
257
258
    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,
Paul McCarthy's avatar
Paul McCarthy committed
259
    and any other arguments with a value of ``LOAD``).
Paul McCarthy's avatar
Paul McCarthy committed
260
261
262
263
264
265
266
267


    **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
Paul McCarthy's avatar
Paul McCarthy committed
268
269
    special ``LOAD`` value. ``None`` is returned for any ``LOAD`` arguments
    corresponded to output files that were not generated by the function.
Paul McCarthy's avatar
Paul McCarthy committed
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
303
304
305


    **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)
Paul McCarthy's avatar
Paul McCarthy committed
306
        atoc = concat('atob.txt', 'btoc.txt', LOAD)[1]
Paul McCarthy's avatar
Paul McCarthy committed
307
308
309
310

        # The inputs are saved to temporary
        # files, and those file names are
        # passed to the concat function.
Paul McCarthy's avatar
Paul McCarthy committed
311
        atoc = concat(np.diag([2, 2, 2, 0]), np.diag([3, 3, 3, 3]), LOAD)[1]
312
313
314
    """


Paul McCarthy's avatar
Paul McCarthy committed
315
316
317
318
319
    def __init__(self, prepareThing, loadThing, *things):
        """Initialise a ``_FileOrThing`` decorator.

        :arg prepareThing: Function which
        :arg loadThing:    Function which is called for arguments that
Paul McCarthy's avatar
Paul McCarthy committed
320
                           were set to :data:`LOAD`.
Paul McCarthy's avatar
Paul McCarthy committed
321
322

        :arg things:
323
        """
Paul McCarthy's avatar
Paul McCarthy committed
324
325
326
        self.__prepareThing = prepareThing
        self.__loadThing    = loadThing
        self.__things       = things
327
328
329


    def __call__(self, func):
Paul McCarthy's avatar
Paul McCarthy committed
330
331
        """Creates and returns the real decorator function. """

332
333
        isFOT   = isinstance(getattr(func, '__self__', None), _FileOrThing)
        wrapper = functools.partial(self.__wrapper, func, isFOT)
Paul McCarthy's avatar
Paul McCarthy committed
334

335
336
337
        # TODO
        wrapper = _update_wrapper(wrapper, func)
        wrapper.__self__ = self
Paul McCarthy's avatar
Paul McCarthy committed
338

339
        return wrapper
Paul McCarthy's avatar
Paul McCarthy committed
340

341

342
    def __wrapper(self, func, isFileOrThing, *args, **kwargs):
Paul McCarthy's avatar
Paul McCarthy committed
343
344
        """Function which wraps ``func``, ensuring that any arguments of
        type ``Thing`` are saved to temporary files, and any arguments
Paul McCarthy's avatar
Paul McCarthy committed
345
        with the value :data:`LOAD` are loaded and returned.
346

347
348
349
350
351
352
353
        :arg func:          The func being wrapped.

        :arg isFileOrThing: Set to ``True`` if ``func`` is a wrapper metho
                            of another ``_FileOrThing`` instance. In this case,
                            the output arguments will be flattenedinto a single
                            tuple.
        """
Paul McCarthy's avatar
Paul McCarthy committed
354
355

        kwargs = kwargs.copy()
356
357
358
        kwargs.update(argsToKwargs(func, args))

        # Create a tempdir to store any temporary
Paul McCarthy's avatar
Paul McCarthy committed
359
        # input/output things, but don't change
360
361
362
363
        # 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
364
            kwargs, infiles, outfiles = self.__prepareThings(td, kwargs)
365
366

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

Paul McCarthy's avatar
Paul McCarthy committed
369
370
            # Load the output things that
            outthings = []
371
            for of in outfiles:
Paul McCarthy's avatar
Paul McCarthy committed
372
            # were specified as LOAD
373
374
375

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

Paul McCarthy's avatar
Paul McCarthy committed
378
                # load the thing
379
                else:
Paul McCarthy's avatar
Paul McCarthy committed
380
                    ot = self.__loadThing(of)
381

Paul McCarthy's avatar
Paul McCarthy committed
382
                outthings.append(ot)
383

Paul McCarthy's avatar
Paul McCarthy committed
384
385
386
            if isFileOrThing:
                things = result[1:]
                result = result[0]
387
                return tuple([result] + list(things) + outthings)
Paul McCarthy's avatar
Paul McCarthy committed
388
389
            else:
                return tuple([result] + outthings)
390
391


Paul McCarthy's avatar
Paul McCarthy committed
392
    def __prepareThings(self, workdir, kwargs):
393
394
395
396
397
398
399
        """
        """

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

Paul McCarthy's avatar
Paul McCarthy committed
400
        for tname in self.__things:
401

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

Paul McCarthy's avatar
Paul McCarthy committed
404
            if tval is None:
405
406
                continue

Paul McCarthy's avatar
Paul McCarthy committed
407
408
409
410
411
412
            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
413
414
415
416

        return kwargs, infiles, outfiles


Paul McCarthy's avatar
Paul McCarthy committed
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
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.
Paul McCarthy's avatar
Paul McCarthy committed
458
        elif val == LOAD:
Paul McCarthy's avatar
Paul McCarthy committed
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
495
496
497
498
499
            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.
Paul McCarthy's avatar
Paul McCarthy committed
500
        elif val == LOAD:
Paul McCarthy's avatar
Paul McCarthy committed
501
502
503
504
505
506
            newval  = op.join(workdir, '{}.txt'.format(name))
            outfile = newval

        return newval, infile, outfile

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