Extending @classmethod

[JoanPuig]

Hi,

I have a custom decorator that looks like this:

import inspect
from functools import wraps
from typing import Type

import attrs


def my_validation(func):
    @wraps(func)
    def wrapper(cls, *args, **kwargs):
        return func(cls, *args, **kwargs)

    wrapper._is_my_validation = True
    return classmethod(wrapper)


def my_function(cls: Type):
    print(f"my_function cls type: {type(cls)}")


@attrs.define
class MyClass:
    @my_validation
    def validate_field_metadata(cls):
        my_function(cls)


for name, method in inspect.getmembers(MyClass, inspect.ismethod):
    if hasattr(method, "_is_my_validation"):
        print(f"is '{name}' a validation method: {method._is_my_validation}")
        method()

The aim is to mark certain class methods within a class as being validation methods. When I run pyrefly on it I get the following errors:

ERROR Object of class `_Wrapped` has no attribute `_is_my_validation` [missing-attribute]
  --> scratch_class_method.py:13:5
   |
13 |     wrapper._is_my_validation = True
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^
   |
ERROR Argument `Self@MyClass` is not assignable to parameter `cls` with type `type[Unknown]` in function `my_function` [bad-argument-type]
  --> scratch_class_method.py:25:21
   |
25 |         my_function(cls)
   |                     ^^^
   |
 INFO 2 errors

What would be the correct way to write the annotation so that pyrefly understands the intent?

Thanks,
Joan

[yangdanny97]

This one seems like a tricky case to handle, it uses dynamic features in Python that are tricky to model statically. We don't have a way to express "this decorator adds an attribute" in the Python type system.

cc @stroxler or @grievejia for opinions here

[stroxler]

There's not a good way to do this and get all the behaviors you might want in Pyrefly today; supporting it is complicated and touches on underspecified behavior (it actually might be a struggle with any type checker today).

Let me first address the question of adding an attribute to a function: the solution there is generally to represent the extra attribute in terms of a callback protocol and make the decorator return it:

class FuncWithValidation[**P, R](Protocol):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: ...
    _is_my_validation: Literal[True]

def my_validation[**P, R](func: Callable[P, R]) -> FuncWithValidation[P, R]:
    @wraps(func)
    def wrapper(cls, *args, **kwargs):
        return func(cls, *args, **kwargs)

    wrapper._is_my_validation = True
    return classmethod(wrapper)

This will work great for static functions not in classes.

But there's a problem: not all callable types implement the descriptor behavior for method binding - for example some builtin functions like abs don't, the result of @staticmethod doesn't, and the result of @classmethod creates a different descriptor.

So what will type checkers assume about FuncWithValidation? Today, I believe most of them assume it has no descriptor behavior at all, with is definitely wrong for most actual examples (in this case the result of @wraps applied to wrapper is a FunctionType with normal behavior, and then we call classmethod on it).

One way around this is to embrace that incorrect analysis by writing the type as:

def my_validation(func: Callable[Concatenate[Any, P], R]) -> FuncWithValidation[P, R]:

so that the decorator actually strips the first argument. This doesn't accurately represent what's actually going on, but it does produce the right final signature; that's actually a common approach people use with mypy today.

One thing you'll run into with that is that Pyrefly will currently guess the wrong type for cls: we don't currently have a good way of knowing the right type so we use the decorators to guess heuristically, and in the absence of a @classmethod decorator we will guess Self instead of type[Self], so you'll need an explicit annotation.

[stroxler]

This is definitely something we'd like to improve in Pyrefly, but unfortunately there are too many constraints and open questions about what the best behavior would be to really begin right now.

[JoanPuig]

HI @yangdanny97 and @stroxler thanks a lot for the detailed explanation. I understand that this is a corner case and it will take some time to decide how to handle. I will keep an eye for future updates.