def _generate_stubs()

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))