Source code for sunkit_image.utils.decorators
import inspect
from functools import wraps
from collections.abc import Callable
import numpy as np
from sunpy.map import GenericMap, Map
__all__ = ["accept_array_or_map"]
[docs]
def accept_array_or_map(*, arg_name: str, output_to_map=True) -> Callable[[Callable], Callable]:
"""
Decorator that allows a function to accept an array or a
`sunpy.map.GenericMap` as an argument.
This can be applied to functions that:
- Take a single array or map as input
- Return a single array that has the same pixel coordinates
as the input array.
Parameters
----------
arg_name : `str`
Name of data/map argument in function signature.
output_to_map : `bool`, optional
If `True` (the default), convert the function return to a map if a map
is given as input. For this to work the decorated function must return
an array where pixels have the same coordinates as the input map data.
"""
def decorate(f: Callable) -> Callable:
sig = inspect.signature(f)
if arg_name not in sig.parameters:
msg = f"Could not find '{arg_name}' in function signature"
raise RuntimeError(msg)
@wraps(f)
def inner(*args, **kwargs) -> np.ndarray | GenericMap:
sig_bound = sig.bind(*args, **kwargs)
map_arg = sig_bound.arguments[arg_name]
if isinstance(map_arg, GenericMap):
map_in = True
sig_bound.arguments[arg_name] = map_arg.data
elif isinstance(map_arg, np.ndarray):
map_in = False
else:
msg = f"'{arg_name}' argument must be a sunpy map or numpy array (got type {type(map_arg)})"
raise TypeError(msg)
# Run decorated function
array_out = f(*sig_bound.args, **sig_bound.kwargs)
return Map(array_out, map_arg.meta) if map_in and output_to_map else array_out
return inner
return decorate