[plugins/help] Triggering node visits?

[eliegoudout]

Hello there 👋

First of all, I post here instead of on Gitter because I think my question is a bit long and it would pollute the thread while being hard to respond too. Also, Thanks for all the documentation in source code, it helps a lot already!

I'm trying to write a plugin for my paramclasses library and I would love some help! For a starting point, I'll focus on my protected function/decorator, which essentially provides runtime equivalent for both typing.finaland typing.Final during paramclasses' definition:

class Foo(ParamClass):
    x = protected(0)  # Unmodifiable in instances and subclasses

    @protected
    def bar(self) -> None:
        """Works on methods."""

    @protected
    @property
    def baz(self) -> int:
        """Also works on properties."""
        return 0

I started understanding a bit more how mypy plugins work, and I currently have the following:

METAPARAMCLASS_FULLNAME = "paramclasses.paramclasses._MetaParamClass"  # Paramclasses have this as metaclass

class CustomPlugin(Plugin):
    """Help mypy undestand ``@protected`` and ``protected()``."""

    def get_metaclass_hook(self, fullname) -> Callable[[ClassDefContext], None] | None:
        """Update paramclass definition."""
        return paramclass_finder_hook

    def get_base_class_hook(self, fullname) -> Callable[[ClassDefContext], None] | None:
        """Update paramclass definition."""
        return paramclass_finder_hook


def paramclass_finder_hook(ctx: ClassDefContext):
    """Identify paramclasses and trigger class def modification."""
    mcs = ctx.cls.info.metaclass_type
    if mcs is not None and mcs.type.fullname == METAPARAMCLASS_FULLNAME:
        modify_paramclass_def(ctx.cls)


def modify_paramclass_def(cls: ClassDef) -> None:
    """Handle ``@property`` and ``protected()``."""
    replace_protected_decorator_with_final(cls)
    replace_protected_assignment_with_Final(cls)


def replace_protected_decorator_with_final(cls: ClassDef) -> None:
    """Visit decorator nodes, move @protected, add to original_decorators, mark node as final."""
    # Help pls!

As you can see, I think I'm pretty close to a first POC, I just need to know how to visit cls in replace_protected_decorator_with_final and modify targetted nodes. For that, I'll take inspiration from mypy.semanal.SemanticAnalyzer.visit_decorator() I think?

But I have so many questions:

1. Is my overall approach sound/in the spirit of mypy plugins API?
2. How can I trigger a visit? I thought about writing a custom semantic analyzer but I don't know how to "trigger" it?
3. If I write one, should my SemanticAnalyzer inherit from (NodeVisitor[None], SemanticAnalyzerInterface, SemanticAnalyzerPluginInterface) just like the one from mypy/semanal.py? I don't really understand what's useful and when.
4. Should I carry out both replace_protected_decorator_with_final and replace_protected_assignment_with_Final simultaneously with a unique semantic analyzer?

Thanks a lot in advance!!!
Élie

[sterliakov]

Fortunately, you do not need any nesting, so a complete visitor is probably an overkill. You can just iterate over class attributes, check if they have the desired decorator (or refer to a call) and mark corresponding nodes as final.

Something like

from mypy.nodes import AssignmentStmt, CallExpr, ClassDef, Decorator, MemberExpr, NameExpr, Var
from mypy.plugin import ClassDefContext, Plugin

PROTECTED_FULLNAME = "paramclasses.paramclasses.protected"

def modify_paramclass_def(cls: ClassDef) -> None:
    """Handle ``@property`` and ``protected()``."""
    for item in cls.defs.body:
        if isinstance(item, Decorator) and any(
            is_protected_call(deco) for deco in item.decorators
        ):
            # This one is easy
            item.var.is_final = True
            item.func.is_final = True
            # Do this if you want this decorator to be otherwise opaque to the type checker
            item.decorators = [
                deco for deco in item.decorators if not is_protected_call(deco)
            ]
        elif (
            isinstance(item, AssignmentStmt)
            and isinstance(item.rvalue, CallExpr)
            and is_protected_call(item.rvalue.callee)
        ):
            # TODO: check that it has exactly one lvalue, that call has exactly one arg,
            # maybe take explicit annotations into account
            item.is_final_def = True
            if isinstance(node := item.lvalues[0].node, Var):
                # actual process is more complex (see `store_final_status` in semanal.py)
                # but this should be enough
                node.is_final = True
            # you can do this, or just define `@protected` in stub as an identity fn
            # I'm not sure that such modifications are safe in incremental mode
            # Uncomment the line below to get rid of _ProtectedType
            # item.rvalue = item.rvalue.args[0]


def is_protected_call(expr):
    return (
        isinstance(expr, NameExpr) and expr.fullname == PROTECTED_FULLNAME
        or isinstance(expr, MemberExpr) and expr.fullname == PROTECTED_FULLNAME
    )

It is probably impossible to get rid of "Decorators on top of @Property are not supported", that is emitted during early semanal, earlier than this hook AFAIC. You can recommend disabling prop-decorator error code if this becomes a problem, or using ignore comments if there only a handful of such @protected @property occurrences.

[eliegoudout]

@sterliakov Thank you sooo much for your answer and the effort you made to provide a functional exemple 🙏

I will try as soon as I can (during next week hopefully) and come back to you!