Is there a way to annotate a function as a type when called at module level but show its property value when inside a class?

Question

While working on a proyect, I stumbled uppon a problem regarding typing annotations when returning custom properties objects.

This is a code snippet that represents the problem:

from typing import TypeVar, Any

_B = TypeVar("_B")

class CustomProperty(property):
    def a_function(): ...

def custom_property(type: _B) -> _B:
    return CustomProperty(type)

custom_property(int).a_function()  # As passed int, typing shows it as an int when I expect CustomProperty

class ClassWithCustomProperty:

    property_holder = custom_property(int)  # This shows as an int, correct

In the code above, custom_property returns an instance of CustomProperty, custom_property can be used either inside a class (Where it would act as a property) and at module level (Where it would be used as a CustomProperty instance)

• How should I annotate the custom_property function to make it show the type CustomProperty when called at module level and to show int when inside a class? is this even possible?

In this proyect using the custom_property/CustomProperty as decorators are not valid options.

I've tried with Union[_B, CustomProperty] but got type[int] | CustomProperty which is not ideal because I expected CustomProperty when called at module level and type[int] when called inside a class

Thank you in advanced, feel free to add comments if anything is unclear or need more context.

Answer 1

If I understand you correctly, this would be the desired result:

module_level = custom_property(int)
reveal_type(module_level)                               # CustomProperty[int]
reveal_type(module_level.a_function())                  # int

class ClassWithCustomProperty:
    property_holder = custom_property(int)
    reveal_type(property_holder)                        # CustomProperty[int]

reveal_type(ClassWithCustomProperty.property_holder)    # CustomProperty[int]
reveal_type(ClassWithCustomProperty().property_holder)  # int

To do so, make CustomProperty generic and overload .__get__():

_{(playgrounds: Mypy, Pyright)}

class CustomProperty[T](property):
    @overload
    def __get__(self, instance: None, owner: type[Any] | None, /) -> Self: ...

    @overload
    def __get__(self, instance: Any, owner: type[Any] | None, /) -> T: ...

    def __get__(self, instance: Any, owner: type[Any] | None = None) -> T | Self:
        return super().__get__(instance, owner)

    def a_function(self) -> T: ...

def custom_property[T](type: type[T]) -> CustomProperty[T]:
    return CustomProperty(type)

instance is None either when the descriptor (i.e. CustomProperty) is accessed from the class (C.property_holder) or when there is no such class/instance at all (module_level), in which case we return the descriptor object (i.e. CustomProperty[T]/Self).

Otherwise, the descriptor is accessed from an instance (C().property_holder) and therefore we return the enclosed value (i.e. T).