How to use return type changed in decorator?

[MisLink]

I have a function like this:

@change_return_annotation(str)
def my_function(x: int, y: int) -> int:
    return x + y

The decorator change the return annotation:

def change_return_annotation(new_return_annotation):
    def decorator(func):
        func.__annotations__['return'] = new_return_annotation
        @functools.wrap(func)
        def wrapper(*args, **kwargs):
            return str(func((args, **kwargs))
        return func
    return decorator

I'd like to show the changed return type in mkdocs. It seems that I need to use InspectorExtension to let griffe import the code. Then I try to create one:

class Inspect(InspectorExtension):
    def inspect_function(self, node: ObjectNode) -> None:
        function: Function = self.inspector.current.members[node.name]  # type: ignore
        for decorator in function.decorators:
            ret = node.obj.__annotations__["return"]
            # not work
            # function.returns = ret


class HybridExtension(hybrid.HybridExtension):
    def __init__(self) -> None:
        super().__init__([Inspect()])


Extension = HybridExtension

Here I'm stuck. How to get an AST node for ret, or do I totally do the wrong thing?

[pawamoy]

Hi @MisLink, I've moved this to a discussion 🙂

I see that you tried to use the hybrid extension, nice work! I will push docs on how to use it soon, but you managed to get quite close.

You could indeed use an inspector extension, but in fact it's just as easy and much more performant to use a visitor extension. Let me show you with an example similar to the one found in this comment:

# changed_return_extension.py
from __future__ import annotations

import ast
import re

from griffe.extensions import VisitorExtension
from griffe.agents.nodes import safe_get_annotation


class ChangedReturnTypeExtension(VisitorExtension):
    def visit_functiondef(self, node: ast.FunctionDef) -> None:
        current = self.visitor.current.members[node.name]  # type: ignore[assignment]
        for decorator_node in node.decorator_list:
            if isinstance(decorator_node, ast.Call) and decorator_node.func.id == "change_return_annotation":
                returns = safe_get_annotation(decorator_node.args[0], parent=current)
                if current.is_function:  # function/method
                    current.returns = returns
                elif current.is_attribute:  # property
                    current.annotation = returns


Extension = ChangedReturnTypeExtension

[MisLink]

Thanks for your reply! I don’t know why, but I didn’t receive any notification from GitHub. It looks much better. I will try in a while.

[pawamoy]

You can also get inspiration from this full example using the new extension system: https://mkdocstrings.github.io/griffe/extensions/#full-example. Or read through the page to learn how this system works and how to do that statically instead of importing the object! Let me know if anything's unclear 🙂