RBerga06 / python-utils

Python Utilities
https://rberga06.github.io/python-utils/
GNU Affero General Public License v3.0
1 stars 0 forks source link

Add a way to safely attach information to a function. #38

Closed RBerga06 closed 1 year ago

RBerga06 commented 1 year ago

Is your feature request related to a problem? Please describe. Currently, there is no obvious and safe way to attach information to a function. The decorator API is to be used as follows:

def count_calls(counter: Mut[int]):
    """Count a function's calls."""
    @decorator(data=lambda: counter)
    def count_calls(__data__, __decorated__, *args, **kwargs):
        # It's actually more complicated, because of exception handling,
        #   but it's not relevant in this example.
        __data__._ += 1
        return __decorated__(*args, **kwargs)
    return count_calls

foo_counter = Mut(0)
@count_calls(foo_counter)
def foo(name: str) -> None:
    print(f"Hello, {name}!")

print(foo_counter)  # 0
foo("World")        # Hello, World!
print(foo_counter)  # 1

Describe the solution you'd like It might be a good idea to allow something like:

@decorator(data=lambda: Mut(0), attr="calls_count")
def count_calls(__data__, __decorated__, *args, **kwargs):
    __data__._ += 1  # Since it's a Mut(...) object,
                     # its value will be updated everywhere,
                     # most notably in `__decorated__.calls_count`.
    return __decorated__(*args, **kwargs)

@count_calls
def foo(name: str) -> None:
    print(f"Hello, {name}!")

foo_counter = getattr(foo, "calls_count")  # or some other, safer, API
print(foo_counter)  # 0
foo("World")        # Hello, World!
print(foo_counter)  # 1

Describe alternatives you've considered N/A

Additional context This feature seems required by #37.

RBerga06 commented 1 year ago

The first problem that comes to my mind is this:

@some_other_decorator  # returns a wrapper to 'foo'
@count_calls  # sets 'calls_count' attribute
def foo(name: str) -> None:
    print(f"Hello, {name}!")

foo_counter = getattr(foo, "calls_count")  # AttributeError, since "calls_count" is defined on foo.__wrapped__ and not on foo itself
RBerga06 commented 1 year ago

The first problem that comes to my mind is this:

@some_other_decorator  # returns a wrapper to 'foo'
@count_calls  # sets 'calls_count' attribute
def foo(name: str) -> None:
    print(f"Hello, {name}!")

foo_counter = getattr(foo, "calls_count")  # AttributeError, since "calls_count" is defined on foo.__wrapped__ and not on foo itself

I cannot reproduce this on Python 3.11. In fact, functools.wraps seems to also copy __dict__, allowing for transparent attribute access.

RBerga06 commented 1 year ago

I believe decorator should define count_calls almost like this (simplified):

def count_calls[F](f: F) -> F:
    __data__ = Mut(0)
    @wraps(f)
    def wrapper(*args, **kwargs):
        __data__._ += 1
        return f(*args, **kwargs)
    setattr(wrapper, "calls_count", __data__)
    return cast(F, wrapper)
RBerga06 commented 1 year ago

Now, we have to design a consistent way of accessing the data. What comes to my mind is something like this:

foo_counter = count_calls.get(foo)

The problem is that it has to be type-checked... We might achieve that by defining count_calls as a class, but that comes with great work and an ugly class-in-function definition.

RBerga06 commented 1 year ago

We might instead define the data attribute as a class, with a decorator staticmethod that should work on the function. For example:

@final
class count_calls(Mut[int]):
    _ATTR: ClassVar[str] = "count_calls"

    def __init__(self):
        super().__init__(0)

    @staticmethod
    @decorator(data=count_calls, attr=count_calls._ATTR)
    def decorate(__data__: count_calls, __decorated__, *args, **kwargs) -> Any:
        __data__._ += 1
        return __decorated__(*args, **kwargs)

    @staticmethod
    def get(func) -> count_calls:
        return getattr(func, count_calls._ATTR)

@count_calls.decorate
def foo(name: str) -> None:
    print(f"Hello, {name}!")

count_calls.get(foo)
RBerga06 commented 1 year ago

Ok, this already looks better. We now obviously need a base class to make this easier. This way we can define count_calls like this:

class count_calls(DecWithDataAttr[Mut[int]]):  # or whatever
    ATTR: Final[str] = "count_calls"  # inferred as ClassVar, as for PEP 591

    def init_data(self) -> Mut[int]:
        """Initialize data"""
        return Mut(0)

    @staticmethod
    def decorator_spec(__data__: Mut[int], __decorated__, *args, **kwargs) -> Any:
        __data__._ += 1
        return __decorated__(*args, **kwargs)

@count_calls()
def foo(name: str) -> None:
    print(f"Hello, {name}!")

counter: Mut[int] = count_calls.get(foo)
RBerga06 commented 1 year ago

It looks like we could define a class-API for decorator definition:

class mydecorator(Decorator[_F]):
    """My decorator."""

    @staticmethod
    def spec(__self__: Self, __decorated__: _F, *args, **kwargs) -> Any:
        return __decorated__(*args, **kwargs)

    @override
    def decorate(f: _F) -> _F:
        # optional
        ... # Modify 'f' as you like
        f = super().decorate(f)  # This applies `self.spec(...)` behaviour
        ... # Modify 'f' as you like
        return f

# To be used like this:
@mydecorator()
def foo(name: str) -> None:
    print(f"Hello, {name}!")

And then DecWithDataAttr might be a subclass of Decorator. I'm opening a separate issue to add this functionality.