Source code for astropy.nddata.decorators

# Licensed under a 3-clause BSD style license - see LICENSE.rst

import warnings
from copy import deepcopy
from functools import wraps
from inspect import signature
from itertools import islice

from astropy.utils.exceptions import AstropyUserWarning

from .nddata import NDData

__all__ = ["support_nddata"]

# All supported properties are optional except "data" which is mandatory!
SUPPORTED_PROPERTIES = ["data", "uncertainty", "mask", "meta", "unit", "wcs", "flags"]

[docs] def support_nddata( _func=None, accepts=NDData, repack=False, returns=None, keeps=None, **attribute_argument_mapping, ): """Decorator to wrap functions that could accept an NDData instance with its properties passed as function arguments. Parameters ---------- _func : callable, None, optional The function to decorate or ``None`` if used as factory. The first positional argument should be ``data`` and take a numpy array. It is possible to overwrite the name, see ``attribute_argument_mapping`` argument. Default is ``None``. accepts : class, optional The class or subclass of ``NDData`` that should be unpacked before calling the function. Default is ``NDData`` repack : bool, optional Should be ``True`` if the return should be converted to the input class again after the wrapped function call. Default is ``False``. .. note:: Must be ``True`` if either one of ``returns`` or ``keeps`` is specified. returns : iterable, None, optional An iterable containing strings which returned value should be set on the class. For example if a function returns data and mask, this should be ``['data', 'mask']``. If ``None`` assume the function only returns one argument: ``'data'``. Default is ``None``. .. note:: Must be ``None`` if ``repack=False``. keeps : iterable. None, optional An iterable containing strings that indicate which values should be copied from the original input to the returned class. If ``None`` assume that no attributes are copied. Default is ``None``. .. note:: Must be ``None`` if ``repack=False``. attribute_argument_mapping : Keyword parameters that optionally indicate which function argument should be interpreted as which attribute on the input. By default it assumes the function takes a ``data`` argument as first argument, but if the first argument is called ``input`` one should pass ``support_nddata(..., data='input')`` to the function. Returns ------- decorator_factory or decorated_function : callable If ``_func=None`` this returns a decorator, otherwise it returns the decorated ``_func``. Notes ----- If properties of ``NDData`` are set but have no corresponding function argument a Warning is shown. If a property is set of the ``NDData`` are set and an explicit argument is given, the explicitly given argument is used and a Warning is shown. The supported properties are: - ``mask`` - ``unit`` - ``wcs`` - ``meta`` - ``uncertainty`` - ``flags`` Examples -------- This function takes a Numpy array for the data, and some WCS information with the ``wcs`` keyword argument:: def downsample(data, wcs=None): # downsample data and optionally WCS here pass However, you might have an NDData instance that has the ``wcs`` property set and you would like to be able to call the function with ``downsample(my_nddata)`` and have the WCS information, if present, automatically be passed to the ``wcs`` keyword argument. This decorator can be used to make this possible:: @support_nddata def downsample(data, wcs=None): # downsample data and optionally WCS here pass This function can now either be called as before, specifying the data and WCS separately, or an NDData instance can be passed to the ``data`` argument. """ if (returns is not None or keeps is not None) and not repack: raise ValueError("returns or keeps should only be set if repack=True.") elif returns is None and repack: raise ValueError("returns should be set if repack=True.") else: # Use empty lists for returns and keeps so we don't need to check # if any of those is None later on. if returns is None: returns = [] if keeps is None: keeps = [] # Short version to avoid the long variable name later. attr_arg_map = attribute_argument_mapping if any(keep in returns for keep in keeps): raise ValueError("cannot specify the same attribute in `returns` and `keeps`.") all_returns = returns + keeps def support_nddata_decorator(func): # Find out args and kwargs func_args, func_kwargs = [], [] sig = signature(func).parameters for param_name, param in sig.items(): if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): raise ValueError("func may not have *args or **kwargs.") try: if param.default == param.empty: func_args.append(param_name) else: func_kwargs.append(param_name) # The comparison to param.empty may fail if the default is a # numpy array or something similar. So if the comparison fails then # it's quite obvious that there was a default and it should be # appended to the "func_kwargs". except ValueError as exc: if ( "The truth value of an array with more than one element " "is ambiguous." in str(exc) ): func_kwargs.append(param_name) else: raise # First argument should be data if not func_args or func_args[0] != attr_arg_map.get("data", "data"): raise ValueError( "Can only wrap functions whose first positional " "argument is `{}`" "".format(attr_arg_map.get("data", "data")) ) @wraps(func) def wrapper(data, *args, **kwargs): bound_args = signature(func).bind(data, *args, **kwargs) unpack = isinstance(data, accepts) input_data = data ignored = [] if not unpack and isinstance(data, NDData): raise TypeError( f"Only NDData sub-classes that inherit from {accepts.__name__}" " can be used by this function" "" ) # If data is an NDData instance, we can try and find properties # that can be passed as kwargs. if unpack: # We loop over a list of pre-defined properties for prop in islice(SUPPORTED_PROPERTIES, 1, None): # We only need to do something if the property exists on # the NDData object try: value = getattr(data, prop) except AttributeError: continue # Skip if the property exists but is None or empty. if prop == "meta" and not value: continue elif value is None: continue # Warn if the property is set but not used by the function. propmatch = attr_arg_map.get(prop, prop) if propmatch not in func_kwargs: ignored.append(prop) continue # Check if the property was explicitly given and issue a # Warning if it is. if propmatch in bound_args.arguments: # If it's in the func_args it's trivial but if it was # in the func_kwargs we need to compare it to the # default. # Comparison to the default is done by comparing their # identity, this works because defaults in function # signatures are only created once and always reference # the same item. # FIXME: Python interns some values, for example the # integers from -5 to 255 (any maybe some other types # as well). In that case the default is # indistinguishable from an explicitly passed kwarg # and it won't notice that and use the attribute of the # NDData. if propmatch in func_args or ( propmatch in func_kwargs and ( bound_args.arguments[propmatch] is not sig[propmatch].default ) ): warnings.warn( "Property {} has been passed explicitly and " "as an NDData property{}, using explicitly " "specified value" "".format( propmatch, "" if prop == propmatch else " " + prop ), AstropyUserWarning, ) continue # Otherwise use the property as input for the function. kwargs[propmatch] = value # Finally, replace data by the data attribute data = if ignored: warnings.warn( "The following attributes were set on the " "data object, but will be ignored by the " "function: " + ", ".join(ignored), AstropyUserWarning, ) result = func(data, *args, **kwargs) if unpack and repack: # If there are multiple required returned arguments make sure # the result is a tuple (because we don't want to unpack # numpy arrays or compare their length, never!) and has the # same length. if len(returns) > 1: if not isinstance(result, tuple) or len(returns) != len(result): raise ValueError( "Function did not return the expected number of arguments." ) elif len(returns) == 1: result = [result] if keeps is not None: for keep in keeps: result.append(deepcopy(getattr(input_data, keep))) resultdata = result[all_returns.index("data")] resultkwargs = { ret: res for ret, res in zip(all_returns, result) if ret != "data" } return input_data.__class__(resultdata, **resultkwargs) else: return result return wrapper # If _func is set, this means that the decorator was used without # parameters so we have to return the result of the # support_nddata_decorator decorator rather than the decorator itself if _func is not None: return support_nddata_decorator(_func) else: return support_nddata_decorator