def test_make_function():
"""Test :func:`~utilipy.utils.functools.make_function`."""
# test function, to copy
def _test_func(x=None):
return x, 2
# /def
sig = inspect.signature(_test_func)
params = list(sig.parameters.values())
params[0] = params[0].replace(default=2)
sig = sig.replace(parameters=params)
test_func = functools.make_function(
_test_func.__code__,
globals_=globals(),
name="made_func",
signature=sig,
)
assert test_func.__name__ == "made_func"
assert _test_func() == (None, 2)
assert test_func() == (2, 2)
assert test_func(1) == (1, 2)
def test_drop_parameter():
"""Test drop_parameter."""
signature = inspect.signature(NS.test_func)
# don't drop anything
newsignature = inspect.drop_parameter(signature, None)
assert newsignature == signature
# drop x, by name
signature = inspect.drop_parameter(signature, "x")
assert list(signature.parameters.values())[0].name == "y"
# drop y, by index
signature = inspect.drop_parameter(signature, 0)
assert list(signature.parameters.values())[0].name == "a"
# drop a, by parameter
signature = inspect.drop_parameter(signature,
list(signature.parameters.values())[0])
assert list(signature.parameters.values())[0].name == "b"
# exception
with pytest.raises(TypeError): # can't drop a set
inspect.drop_parameter(signature, set())
return
def test_modify_parameter():
"""Test modify_parameter."""
signature = inspect.signature(NS.test_func)
# ----------------------------------------------------
# doing piecemeal for full code coverage
# and so can cycle through index / versus name as `param` arg
sig = inspect.modify_parameter(signature, "x", name="xx")
sig = inspect.modify_parameter(sig, "xx", kind=inspect.POSITIONAL_ONLY)
sig = inspect.modify_parameter(sig, "y", default="Y")
sig = inspect.modify_parameter(sig, 1, annotation="YY")
parameters = list(sig.parameters.values())
# changed
assert parameters[0].name == "xx"
assert parameters[0].kind == inspect.POSITIONAL_ONLY
assert parameters[1].default == "Y"
assert parameters[1].annotation == "YY"
# not changed
assert parameters[0].annotation == int
assert parameters[1].name == "y"
return
def test_get_kwonlydefaults_from_signature():
"""Test get_annotations_from_signature."""
signature = inspect.signature(NS.test_func)
kwdefaults = inspect.get_kwonlydefaults_from_signature(signature)
assert kwdefaults == {"j": "a", "k": "b"}
return
def test_get_defaults_from_signature():
"""Test get_annotations_from_signature."""
signature = inspect.signature(NS.test_func)
defaults = inspect.get_defaults_from_signature(signature)
assert defaults == (1, 2)
return
def test_append_parameter():
"""Test append_parameter."""
signature = inspect.signature(NS.appendable_func)
sig = inspect.modify_parameter(signature, 0, name="xx")
new_param = list(sig.parameters.values())[0]
signature = inspect.append_parameter(signature, new_param)
assert list(signature.parameters.values())[2].name == "xx"
return
def test_insert_parameter():
"""Test insert_parameter."""
signature = inspect.signature(NS.test_func)
sig = inspect.modify_parameter(signature, 0, name="xx")
new_param = list(sig.parameters.values())[0]
signature = inspect.insert_parameter(signature, 1, new_param)
assert list(signature.parameters.values())[1].name == "xx"
return
def from_amuse_decorator(
function: T.Callable = None, *, arguments: list = []
) -> T.Callable:
"""Function decorator to convert inputs to Astropy quantities.
Parameters
----------
function : types.FunctionType or None, optional
the function to be decoratored
if None, then returns decorator to apply.
arguments : list, optional
arguments to convert
integers are indices into `arguments`
strings are names of `kw` arguments
Returns
-------
wrapper : types.FunctionType
wrapper for function
does a few things
includes the original function in a method `.__wrapped__`
"""
from .convert import from_amuse # TODO, to prevent circular import
if not all([isinstance(a, (int, str)) for a in arguments]):
raise TypeError("elements of `arguments` must be int or str")
if function is None: # allowing for optional arguments
return functools.partial(from_amuse_decorator, arguments=arguments)
sig = inspect.signature(function)
pnames = tuple(sig.parameters.keys())
@functools.wraps(function)
def wrapper(*args, **kw):
"""Wrapper docstring."""
ba = sig.bind_partial(*args, **kw)
ba.apply_defaults()
for i in arguments:
if isinstance(i, str):
ba.arguments[i] = from_amuse(ba.arguments[i])
else: # int
ba.arguments[pnames[i]] = from_amuse(ba.arguments[pnames[i]])
return function(*ba.args, **ba.kwargs)
# /def
return wrapper
def test_get_kinds_from_signature():
"""Test get_annotations_from_signature."""
signature = inspect.signature(NS.test_func)
kinds = inspect.get_kinds_from_signature(signature)
assert kinds == (
inspect.POSITIONAL_OR_KEYWORD,
inspect.POSITIONAL_OR_KEYWORD,
inspect.POSITIONAL_OR_KEYWORD,
inspect.POSITIONAL_OR_KEYWORD,
inspect.VAR_POSITIONAL,
inspect.KEYWORD_ONLY,
inspect.KEYWORD_ONLY,
inspect.VAR_KEYWORD,
)
def test_get_annotations_from_signature():
"""Test get_annotations_from_signature."""
signature = inspect.signature(NS.test_func)
annotations = inspect.get_annotations_from_signature(signature)
assert annotations == {
"x": int,
"a": int,
"args": str,
"j": str,
"kw": dict,
"return": bool,
}
return
def test_replace_with_parameter():
"""Test replace_with_parameter."""
signature = inspect.signature(NS.test_func)
sig = inspect.modify_parameter(signature, 0, name="xx")
sig = inspect.modify_parameter(sig, 1, name="yy")
new_param0 = list(sig.parameters.values())[0]
new_param1 = list(sig.parameters.values())[1]
signature = inspect.replace_with_parameter(signature, "x", new_param0)
signature = inspect.replace_with_parameter(signature, 1, new_param1)
parameters = list(sig.parameters.values())
assert parameters[0].name == "xx"
assert parameters[1].name == "yy"
return
def test_copy_function():
"""Test `~utilipy.utils.functools.copy_function`."""
# test function
def test_func(x: float, y, a=2, b=3, *args, p="L", q="M", **kwargs):
return x, y, a, b, args, p, q, kwargs
tfc = functools.copy_function(test_func)
# ------------------------------------
# test properties
sep_tests = ( # require separate tests
"__call__",
"__delattr__",
"__dir__",
"__eq__",
"__format__",
"__ge__",
"__get__",
"__getattribute__",
"__gt__",
"__hash__",
"__init__",
"__le__",
"__lt__",
"__ne__",
"__reduce__",
"__reduce_ex__",
"__repr__",
"__setattr__",
"__sizeof__",
"__str__",
)
for attr in set(dir(test_func)).difference(sep_tests):
assert getattr(tfc, attr) == getattr(test_func, attr), attr
# __call__
ba = inspect.signature(test_func).bind(
0,
[1, 1.5],
2.1,
3.1,
4.0,
4.33,
4.66,
p=5,
q=[6, 6.5],
r=7,
s=[8, 8.5],
)
assert tfc.__call__(*ba.args, **ba.kwargs) == test_func.__call__(
*ba.args, **ba.kwargs)
# __delattr__
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
tfc.attr_to_del = "delete me"
test_func.attr_to_del = "delete me"
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
assert tfc.attr_to_del == test_func.attr_to_del # test equality
tfc.__delattr__("attr_to_del")
test_func.__delattr__("attr_to_del")
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
# __setattr__
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
tfc.__setattr__("attr_to_del", "delete me")
test_func.__setattr__("attr_to_del", "delete me")
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
assert tfc.attr_to_del == test_func.attr_to_del # and be equal
del tfc.attr_to_del, test_func.attr_to_del
assert hasattr(tfc, "attr_to_del") == hasattr(test_func, "attr_to_del")
# __dir__
assert set(tfc.__dir__()) == set(test_func.__dir__())
# __getattribute__
assert set(tfc.__getattribute__("__dir__")()) == set(
test_func.__getattribute__("__dir__")())
# __eq__, __ge__, __gt__, '__le__', '__lt__', '__ne__'
for attr in ("__eq__", "__ge__", "__gt__", "__le__", "__lt__", "__ne__"):
assert tfc.__getattribute__(attr)(tfc) == test_func.__getattribute__(
attr)(test_func)
assert tfc.__getattribute__(attr)(
test_func) == test_func.__getattribute__(attr)(tfc)
# __format__
# TODO
# __get__
# TODO
# __hash__
# TODO
# __init__
# TODO
# __reduce__, __reduce_ex__ TypeError: can't pickle function objects
# __repr__
# TODO, right now they are the same. Problem?
# __sizeof__
assert tfc.__sizeof__() == test_func.__sizeof__()
# __str__
assert tfc.__str__() != test_func.__str__()
assert tfc.__str__().split(" ")[:-1] == test_func.__str__().split(" ")[:-1]
# ------------------------------------
# test calling
# TODO
return
def __call__(self, wrapped_func: T.Callable) -> T.Callable:
"""Wrap function.
Works by making a wrapper which will convert input and
output arguments to the specified data type.
Parameters
----------
wrapped_func : callable
Function to wrap.
"""
sig = inspect.signature(wrapped_func)
@functools.wraps(wrapped_func)
def wrapper(*args: T.Any, **kwargs: T.Any) -> T.Any:
ba = sig.bind_partial(*args, **kwargs)
ba.apply_defaults()
# PRE
# making arguments self._dtype
if self._inargs is None: # no conversion needed
pass
elif isinstance(self._inargs, slice):
# converting inargs to list of indices
lna = len(ba.args)
inkeys = tuple(ba.arguments.keys())[:lna][self._inargs]
inargs = tuple(range(lna))[self._inargs]
# converting to desired dtype
for k, i in zip(inkeys, inargs):
ba.arguments[k] = self._dtype(ba.args[i])
else: # any iterable
lna = len(ba.args)
argkeys = tuple(ba.arguments.keys())
for i in self._inargs:
if isinstance(i, int): # it's for args
ba.arguments[argkeys[i]] = self._dtype(args[i])
else: # isinstance(i, str)
ba.arguments[i] = self._dtype(ba.arguments[i])
# /PRE
return_ = wrapped_func(*ba.args, **ba.kwargs)
# POST
# no conversion needed
if self._outargs is None or not isinstance(return_, tuple):
return return_
# slice
elif isinstance(self._outargs, slice):
iterable = tuple(range(len(return_)))[self._outargs]
else:
iterable = self._outargs
return_ = list(return_)
for i in iterable:
return_[i] = self._dtype(return_[i])
return tuple(return_)
# /POST
# /def
return wrapper
def random_generator_from_seed(
function: T.Callable = None,
seed_names: T.Union[str, T.Sequence[str]] = ("random", "random_seed"),
generator: T.Callable = np.random.RandomState,
raise_if_not_int: bool = False,
):
"""Function decorator to convert random seed to random number generator.
Parameters
----------
function : types.FunctionType or None (optional)
the function to be decoratored
if None, then returns decorator to apply.
seed_names : list (optional)
possible parameter names for the random seed
generator : ClassType (optional)
ex :class:`numpy.random.default_rng`, :class:`numpy.random.RandomState`
raise_if_not_int : bool (optional, keyword-only)
raise TypeError if seed argument is not an int.
Returns
-------
wrapper : types.FunctionType
wrapper for function
converts random seeds to random number generators before calling.
includes the original function in a method `.__wrapped__`
Raises
------
TypeError
If `raise_if_not_int` is True and seed argument is not an int.
"""
if isinstance(seed_names, str): # correct a bare string to list
seed_names = (seed_names,)
if function is None: # allowing for optional arguments
return functools.partial(
random_generator_from_seed,
seed_names=seed_names,
generator=generator,
)
sig = inspect.signature(function)
pnames = tuple(sig.parameters.keys())
@functools.wraps(
function,
_doc_fmt={"seed_names": seed_names, "random_generator": generator},
)
def wrapper(*args, **kw):
"""Wrapper docstring, added to Function.
Notes
-----
T.Any argument in {seed_names} will be interpreted as a random seed,
if it is an integer, and will be converted to a random number generator
of type {random_generator}.
"""
ba = sig.bind_partial(*args, **kw)
ba.apply_defaults()
# go through possible parameter names for the random seed
# if it is a parameter and the value is an int, change to RandomState
for name in seed_names: # iterate through possible
if name in pnames: # see if present
if isinstance(ba.arguments[name], int): # seed -> generator
ba.arguments[name] = generator(ba.arguments[name])
elif raise_if_not_int:
raise TypeError(f"{name} must be <int>")
else: # do not replace
pass
# /for
return function(*ba.args, **ba.kwargs)
# /def
return wrapper
def add_folder_backslash(
function=None,
*,
arguments: T.List[T.Union[str, int]] = [],
_doc_style="numpy",
_doc_fmt={},
):
"""Add backslashes to str arguments.
For use in ensuring directory file-paths end in '/', when
``os.join`` just won't do.
Parameters
----------
function : T.Callable or None, optional
the function to be decoratored
if None, then returns decorator to apply.
arguments : list of string or int, optional
arguments to which to append '/', if not already present
strings are names of arguments.
Can also be int, which only applies to args.
Returns
-------
wrapper : T.Callable
wrapper for function
does a few things
includes the original function in a method ``.__wrapped__``
Other Parameters
----------------
_doc_style: str or formatter, optional
default 'numpy'
parameter to `~utilipy.wraps`
_doc_fmt: dict, optional
default None
parameter to `~utilipy.wraps`
Examples
--------
For modifying a single argument
>>> @add_folder_backslash(arguments='path')
... def func(path):
... return path
>>> func("~/Documents")
'~/Documents/'
When several arguments need modification.
>>> @add_folder_backslash(arguments=('path1', 'path2'))
... def func(path1, path2):
... return (path1, path2)
>>> func("~/Documents", "~Desktop")
('~/Documents/', '~Desktop/')
"""
if isinstance(arguments, (str, int)): # recast as tuple
arguments = (arguments,)
if function is None: # allowing for optional arguments
return functools.partial(
add_folder_backslash,
arguments=arguments,
_doc_style=_doc_style,
_doc_fmt=_doc_fmt,
)
sig = inspect.signature(function)
@functools.wraps(function, _doc_style=_doc_style, _doc_fmt=_doc_fmt)
def wrapper(*args, **kw):
"""Wrapper docstring.
Parameters
----------
store_inputs: bool
whether to store function inputs in a BoundArguments instance
default {store_inputs}
"""
# bind args & kwargs to function
ba = sig.bind_partial(*args, **kw)
ba.apply_defaults()
for name in arguments: # iter through args
# first check it's a string
if not isinstance(ba.arguments[name], (str, bytes)):
continue
else:
str_type = type(ba.arguments[name]) # get string type
backslash = str_type("/") # so we can work with any type
if isinstance(name, int): # only applies to args
if not ba.args[name].endswith(backslash):
ba.args[name] += backslash
elif isinstance(name, str): # args or kwargs
if not ba.arguments[name].endswith(backslash):
ba.arguments[name] += backslash
else:
raise TypeError("elements of `args` must be int or str")
return function(*ba.args, **ba.kwargs)
# /def
return wrapper
def __call__(self, wrapped_function: T.Callable):
"""Make decorator.
Parameters
----------
wrapped_function : Callable
function to wrap
Returns
-------
wrapped: Callable
wrapped function
"""
# Extract the function signature for the function we are wrapping.
wrapped_signature = inspect.signature(wrapped_function)
@functools.wraps(wrapped_function)
def wrapped(
*func_args: T.Any,
unit: UnitableType = self.unit,
to_value: bool = self.to_value,
equivalencies: T.Sequence = self.equivalencies,
decompose: T.Union[bool, T.Sequence] = self.decompose,
assumed_units: dict = self.assumed_units,
_skip_decorator: bool = False,
**func_kwargs: T.Any,
):
# skip the decorator
if _skip_decorator:
return wrapped_function(*func_args, **func_kwargs)
# make func_args editable
_func_args: list = list(func_args)
# Bind the arguments to our new function to the signature of the original.
bound_args = wrapped_signature.bind(*_func_args, **func_kwargs)
# Iterate through the parameters of the original signature
for i, param in enumerate(wrapped_signature.parameters.values()):
# We do not support variable arguments (*args, **kwargs)
if param.kind in {
inspect.Parameter.VAR_KEYWORD,
inspect.Parameter.VAR_POSITIONAL,
}:
continue
# Catch the (never) case where bind relied on a default value.
if (
param.name not in bound_args.arguments
and param.default is not param.empty
):
bound_args.arguments[param.name] = param.default
# Get the value of this parameter (argument to new function)
arg = bound_args.arguments[param.name]
# +----------------------------------+
# Get default unit or physical type,
# either from decorator kwargs
# or annotations
if param.name in assumed_units:
dfunit = assumed_units[param.name]
elif self.assume_annotation_units is True:
dfunit = param.annotation
# elif not assumed_units:
# dfunit = param.annotation
else:
dfunit = inspect.Parameter.empty
adjargbydfunit = True
# If the dfunit is empty, then no target units or physical
# types were specified so we can continue to the next arg
if dfunit is inspect.Parameter.empty:
adjargbydfunit = False
# If the argument value is None, and the default value is None,
# pass through the None even if there is a dfunit unit
elif arg is None and param.default is None:
adjargbydfunit = False
# Here, we check whether multiple dfunit unit/physical type's
# were specified in the decorator/annotation, or whether a
# single string (unit or physical type) or a Unit object was
# specified
elif isinstance(dfunit, str):
dfunit = _get_allowed_units([dfunit])[0]
elif not isiterable(dfunit):
pass
else:
raise ValueError("target must be one Unit, not list")
if (not hasattr(arg, "unit")) & (adjargbydfunit is True):
if i < len(_func_args):
# print(i, len(bound_args.args))
_func_args[i] *= dfunit
else:
func_kwargs[param.name] *= dfunit
arg *= dfunit
# +----------------------------------+
# Get target unit or physical type,
# from decorator kwargs or annotations
if param.name in self.decorator_kwargs:
targets = self.decorator_kwargs[param.name]
else:
targets = param.annotation
# If the targets is empty, then no target units or physical
# types were specified so we can continue to the next arg
if targets is inspect.Parameter.empty:
continue
# If the argument value is None, and the default value is None,
# pass through the None even if there is a target unit
if arg is None and param.default is None:
continue
# Here, we check whether multiple target unit/physical type's
# were specified in the decorator/annotation, or whether a
# single string (unit or physical type) or a Unit object was
# specified
if isinstance(targets, str) or not isiterable(targets):
valid_targets = [targets]
# Check for None in the supplied list of allowed units and, if
# present and the passed value is also None, ignore.
elif None in targets:
if arg is None:
continue
else:
valid_targets = [t for t in targets if t is not None]
if not hasattr(arg, "unit"):
arg = arg * dimensionless_unscaled
valid_targets.append(dimensionless_unscaled)
else:
valid_targets = targets
# Now loop over the allowed units/physical types and validate
# the value of the argument:
_validate_arg_value(
param.name,
wrapped_function.__name__,
arg,
valid_targets,
self.equivalencies,
)
# # evaluated wrapped_function
with add_enabled_equivalencies(equivalencies):
return_ = wrapped_function(*_func_args, **func_kwargs)
# if func_kwargs:
# return_ = wrapped_function(*_func_args, **func_kwargs)
# else:
# return_ = wrapped_function(*_func_args)
if (
wrapped_signature.return_annotation
not in (inspect.Signature.empty, None)
and unit is None
):
unit = wrapped_signature.return_annotation
return quantity_return_(
return_,
unit=unit,
to_value=to_value,
equivalencies=equivalencies,
decompose=decompose,
)
# /def
# TODO dedent
# wrapped.__doc__ = inspect.cleandoc(wrapped.__doc__ or "") + _funcdec
wrapped.__doc__ = wrapped_function.__doc__
return wrapped
