I'm writing a decorator which adds two member functions to each field of a dataclass
:
dataclass_instance.SET_a(value, time)
. Let's assume that the time is represented by an integer.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
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 .