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
223
        kwargs[name] = val

    return kwargs


RETURN = 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
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
303
304
305
306
307
308
309
310
311
312
    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]
313
314
315
    """


Paul McCarthy's avatar
Paul McCarthy committed
316
317
318
319
320
321
322
323
    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:
324
        """
Paul McCarthy's avatar
Paul McCarthy committed
325
326
327
        self.__prepareThing = prepareThing
        self.__loadThing    = loadThing
        self.__things       = things
328
329
330


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

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

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

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

342

343
    def __wrapper(self, func, isFileOrThing, *args, **kwargs):
Paul McCarthy's avatar
Paul McCarthy committed
344
345
346
        """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.
347

348
349
350
351
352
353
354
        :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
355
356

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

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

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

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

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

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

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

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


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

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

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

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

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

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

        return kwargs, infiles, outfiles


Paul McCarthy's avatar
Paul McCarthy committed
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
495
496
497
498
499
500
501
502
503
504
505
506
507
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)