You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
84 lines
2.9 KiB
84 lines
2.9 KiB
"""Useful decorators for Traitlets users."""
|
|
|
|
import copy
|
|
from inspect import Parameter, Signature, signature
|
|
|
|
from ..traitlets import Undefined
|
|
|
|
|
|
def _get_default(value):
|
|
"""Get default argument value, given the trait default value."""
|
|
return Parameter.empty if value == Undefined else value
|
|
|
|
|
|
def signature_has_traits(cls):
|
|
"""Return a decorated class with a constructor signature that contain Trait names as kwargs."""
|
|
traits = [
|
|
(name, _get_default(value.default_value))
|
|
for name, value in cls.class_traits().items()
|
|
if not name.startswith("_")
|
|
]
|
|
|
|
# Taking the __init__ signature, as the cls signature is not initialized yet
|
|
old_signature = signature(cls.__init__)
|
|
old_parameter_names = list(old_signature.parameters)
|
|
|
|
old_positional_parameters = []
|
|
old_var_positional_parameter = None # This won't be None if the old signature contains *args
|
|
old_keyword_only_parameters = []
|
|
old_var_keyword_parameter = None # This won't be None if the old signature contains **kwargs
|
|
|
|
for parameter_name in old_signature.parameters:
|
|
# Copy the parameter
|
|
parameter = copy.copy(old_signature.parameters[parameter_name])
|
|
|
|
if (
|
|
parameter.kind is Parameter.POSITIONAL_ONLY
|
|
or parameter.kind is Parameter.POSITIONAL_OR_KEYWORD
|
|
):
|
|
old_positional_parameters.append(parameter)
|
|
|
|
elif parameter.kind is Parameter.VAR_POSITIONAL:
|
|
old_var_positional_parameter = parameter
|
|
|
|
elif parameter.kind is Parameter.KEYWORD_ONLY:
|
|
old_keyword_only_parameters.append(parameter)
|
|
|
|
elif parameter.kind is Parameter.VAR_KEYWORD:
|
|
old_var_keyword_parameter = parameter
|
|
|
|
# Unfortunately, if the old signature does not contain **kwargs, we can't do anything,
|
|
# because it can't accept traits as keyword arguments
|
|
if old_var_keyword_parameter is None:
|
|
raise RuntimeError(
|
|
"The {} constructor does not take **kwargs, which means that the signature can not be expanded with trait names".format(
|
|
cls
|
|
)
|
|
)
|
|
|
|
new_parameters = []
|
|
|
|
# Append the old positional parameters (except `self` which is the first parameter)
|
|
new_parameters += old_positional_parameters[1:]
|
|
|
|
# Append *args if the old signature had it
|
|
if old_var_positional_parameter is not None:
|
|
new_parameters.append(old_var_positional_parameter)
|
|
|
|
# Append the old keyword only parameters
|
|
new_parameters += old_keyword_only_parameters
|
|
|
|
# Append trait names as keyword only parameters in the signature
|
|
new_parameters += [
|
|
Parameter(name, kind=Parameter.KEYWORD_ONLY, default=default)
|
|
for name, default in traits
|
|
if name not in old_parameter_names
|
|
]
|
|
|
|
# Append **kwargs
|
|
new_parameters.append(old_var_keyword_parameter)
|
|
|
|
cls.__signature__ = Signature(new_parameters)
|
|
|
|
return cls
|