in metaflow/cmd/develop/stub_generator.py [0:0]
def _generate_stubs(self):
for name, attr in self._current_objects.items():
self._current_parent_module = inspect.getmodule(attr)
self._current_name = name
if inspect.isclass(attr):
self._stubs.append(self._generate_class_stub(name, attr))
elif inspect.isfunction(attr):
# Special handling of the `step` function where we want to add an
# overload. This is just a single case so we don't make it general.
# Unfortunately, when iterating, it doesn't see the @overload
if (
name == "step"
and self._current_module_name == self._root_module[:-1]
):
self._stubs.append(
self._generate_function_stub(
name,
func=attr,
sign=[
inspect.Signature(
parameters=[
inspect.Parameter(
name="f",
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=Callable[
[FlowSpecDerived], None
],
)
],
return_annotation=Callable[
[FlowSpecDerived, StepFlag], None
],
),
inspect.Signature(
parameters=[
inspect.Parameter(
name="f",
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=Callable[
[FlowSpecDerived, Any], None
],
)
],
return_annotation=Callable[
[FlowSpecDerived, Any, StepFlag], None
],
),
inspect.signature(attr),
],
)
)
else:
self._stubs.append(self._generate_function_stub(name, attr))
elif isinstance(attr, functools.partial):
if issubclass(attr.args[0], Decorator):
# Special case where we are going to extract the parameters from
# the docstring to make the decorator look nicer
res = self._extract_signature_from_decorator(
name,
attr.args[0].__doc__,
is_flow_decorator=issubclass(attr.args[0], FlowDecorator),
)
if res:
self._stubs.append(
self._generate_function_stub(
name,
func=attr.func,
sign=[r[0] for r in res],
doc=res[-1][1],
)
)
else:
# print(
# "WARNING: Could not extract decorator signature for %s"
# % name
# )
pass
else:
self._stubs.append(
self._generate_function_stub(
name, attr.func, doc=attr.args[0].__doc__
)
)
elif not inspect.ismodule(attr):
self._stubs.append(self._generate_generic_stub(name, attr))