I have a python decorator. All functions it decorates must accept a magic argument. If the magic argument is supplied, the decorated function is evaluated immediately and returned. If the magic argument is not supplied, a partial function is returned instead, allowing for lazy evaluation later.

I am writing a mypy plugin to type the decorated functions. I need to do two things:

1. Change the decorated function's return type to PartialFunction if the magic argument is not supplied.
2. Prevent "Missing positional argument" for magic when a PartialFunction is returned

I can do 1 easily enough with a get_function_hook callback.

However 2 is more difficult. If I use a signature callback (via get_function_signature_hook) it changes the function signature for all invocations (not just the ones without the magic argument). I've tried changing args in a get_function_hook callback (with .copy_modified() and creating new Instance objects) but without success.

Is this possible? And if so, how should I approach this? Examples or links to documentation very welcome!

Below a toy example:

mypy.ini

[mypy]
plugins = partial_plugin

This the mypy plugin: partial_plugin.py

from collections.abc import Callable

from mypy.plugin import FunctionContext, Plugin
from mypy.types import Type


def _partial_function_hook_callback(ctx: FunctionContext) -> Type:
    if "magic" in ctx.callee_arg_names:
        magic_index = ctx.callee_arg_names.index("magic")
        if not ctx.args[magic_index]:  # If magic not supplied, return PartialFunction type
            return ctx.api.named_type("my_partial.PartialFunction")
    return ctx.default_return_type


class PartialFunctionPlugin(Plugin):
    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
        """Return a hook for the given function name if it matches 'my_partial'."""
        if fullname.startswith("my_partial."):
            return _partial_function_hook_callback
        return None


def plugin(version: str) -> type[PartialFunctionPlugin]:
    """Entry point for the plugin."""
    return PartialFunctionPlugin

The decorator and examples to type-check : my_partial.py

import functools
from collections.abc import Callable
from typing import Any, ParamSpec, TypeVar, reveal_type

T = TypeVar("T")
P = ParamSpec("P")


class PartialFunction:
    def __init__(self, func: Callable[P, T], *args: Any, **kwargs: Any) -> None:
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        # Combine the args and kwargs with the stored ones
        combined_args = self.args + args
        combined_kwargs = {**self.kwargs, **kwargs}
        return self.func(*combined_args, **combined_kwargs)


def partial_decorator(func: Callable[P, T]) -> Callable[P, T]:
    def decorator_inner(*args: Any, **kwargs: Any) -> Any:
        if "magic" in kwargs:
            # If the magic argument is passed, evaluate immediately
            return func(*args, **kwargs)
        # Otherwise, we want to return a partial function
        return PartialFunction(func, *args, **kwargs)

    return functools.update_wrapper(decorator_inner, func)


#####################


@partial_decorator
def concat(x: str, magic: str | None) -> str:
    return f"{x} {magic}"


foo = concat("hello", magic="world")  # gives "hello world"
reveal_type(foo)  # Should be type 'str'
print(foo)

foo_partial = concat("hello")  # `magic` not supplied, returns a partial function
reveal_type(foo_partial)  # Should be type 'PartialFunction'

print(foo_partial(magic="everyone"))  # gives "hello everyone"

With the above three files in the same directory I can run it through mypy with python -m mypy my_partial.py

Doing that gives this result:

my_partial.py:42: note: Revealed type is "builtins.str"
my_partial.py:45: error: Missing positional argument "magic" in call to "concat"  [call-arg]
my_partial.py:46: note: Revealed type is "my_partial.PartialFunction"

The two revealed types are good. It is the "error" I want to fix. Thanks!