Search code examples
python-decoratorspython-dataclassespython-3.11

Setting the signature of injected member functions in decorators


I'm writing a decorator which adds two member functions to each field of a dataclass:

  1. The first member function can be used to set the value of the field and the time when the field is updated, e.g., dataclass_instance.SET_a(value, time). Let's assume that the time is represented by an integer.
  2. With the second member function, the time can be get for the field, e.g., dataclass_instance.GET_a().

I created a decorator for this purpose:

def AddLastUpdateTime(cls):
    cls_annotations = cls.__dict__.get('__annotations__', {})

    setattr(cls, 'last_update_time', dict())
    cls_annotations['last_update_time'] = dict[str, int]

    def create_custom_setter(field):
        def custom_setter(self, value: field.type, last_update_time: int) -> None:
            setattr(self, field.name, value)
            self.last_update_time[field.name] = last_update_time
        return custom_setter

    def create_custom_getter(field) -> Callable[[], int]:
        def custom_getter(self) -> int:
            return self.last_update_time[field.name]
        return custom_getter

    for field in fields(cls):
        temp = create_custom_setter(field)
        temp.__doc__ = f'This is a setter for {field.name}.'
        setattr(cls, f'SET_{field.name}', temp)
        cls_annotations[f'SET_{field.name}'] = Callable[[
            field.type, int], None]

        temp = create_custom_getter(field)
        temp.__doc__ = f'This is a getter for {field.name}.'
        getattr(cls, f'GET_{field.name}', temp)
        cls_annotations[f'GET_{field.name}'] = Callable[[], field.type]

    print(cls_annotations)
    return cls

When I use it like this:

@AddLastUpdateTime
@dataclass
class A:
    a: float

it works as expected, but the type hints do not show up properly.

dataclass_instance.SET_a? returns the correct signature in a Jupyter Notebook: Signature: a.SET_a(value: float, last_update_time: int) -> None but the dynamic type hint suggestion in VSCode only shows (function) SET_a: Any

How can I make the decorator add a type hint to the injected member functions correctly? Thanks!

UPDATE: I tried many approaches, none of them seemed to work as expected. Finally, I decided to write a function which generates a new module based on the fields of a dataclass like this:

from __future__ import annotations

from dataclasses import dataclass, fields
from typing import TypeVar, Type
import pathlib

T = TypeVar('T')


def extend_with_time(cls: Type[T]):
    print('Decorator runs')
    path_folder = pathlib.Path(__file__).parent
    path_new_module = path_folder / pathlib.Path(f'_with_time_{cls.__name__}.py')
    path_new_module.unlink(missing_ok=True)

    with path_new_module.open('w') as file:
        file.write(f'from . import main\n')
        file.write(f'from dataclasses import dataclass, field\n')

        file.write(f'print(r"{path_new_module} imported")\n')

        file.write(f'@dataclass\n')
        file.write(f'class WithTime(main.{cls.__name__}):\n')
        file.write('\t_time_dict: dict[str, int] = field(default_factory=lambda: {})\n')
        for field in fields(cls):
            file.write(f'\tdef set_{field.name}(self, value: {field.type}, time: int)->None:\n')
            file.write(f'\t\tself.{field.name} = value\n')
            file.write(f'\t\tself._time_dict["{field.name}"] = time\n')

            file.write(f'\tdef get_time_{field.name}(self)->int:\n')
            file.write(f'\t\treturn self._time_dict.get("{field.name}", 0)\n')

@dataclass
class A:
    a: float = 9
    b: int | None = None

extend_with_time(A)
from ._with_time_A import WithTime as AWithTime


Solution

  • you are creating dynamic functions ( methods created at run time) and expecting technology which uses static code analysis to be aware of that: it won't work, under any approach.

    It works on notebooks, because they actually import the code, and it is "live" with all the dynamic methods already created, and these are introspected.

    The way out, if you really want to bother about this, is to have your decorators check for some especial mark when running (maybe a switch passed at the command line), which triggers then to generate a .pyi file for each file where they are called, and in this .pyi file the methods it creates dynamically are rendered as stubs. You then run your code in this way at test time, and after that pass, static type checking tools can make use of the pyi files. check https://peps.python.org/pep-0484/#stub-files .