diff --git a/pystencils/cache.py b/pystencils/cache.py index b8ac2b06e74815b3edf304f23f5bde1e38dc27bd..d8988f48bc6973f72f987ca361c08b030ce34dec 100644 --- a/pystencils/cache.py +++ b/pystencils/cache.py @@ -1,6 +1,6 @@ import os from collections.abc import Hashable -from functools import partial +from functools import partial, wraps from itertools import chain try: @@ -43,6 +43,7 @@ def sharedmethodcache(cache_id: str): Of course, for this to be useful, said methods must have the same signature (up to additional kwargs) and must return the same result when called with the same arguments.""" def _decorator(user_method): + @wraps(user_method) def _decorated_func(self, *args, **kwargs): objdict = self.__dict__ cache = objdict.setdefault(cache_id, dict()) diff --git a/pystencils_tests/test_sharedmethodcache.py b/pystencils_tests/test_sharedmethodcache.py index 7489dd61bdb8943dbe7fe7998998e002d314cba2..eafb45933cdf0894b91e10db4f3428859986785e 100644 --- a/pystencils_tests/test_sharedmethodcache.py +++ b/pystencils_tests/test_sharedmethodcache.py @@ -60,13 +60,16 @@ class Triad: @sharedmethodcache("triad_cache") def triad(self, a, b, c=0): + """Computes the triad a*b+c.""" self.triad_called += 1 return a * b + c -def test_triab_memoization(): +def test_triad_memoization(): triad = Triad() + assert triad.triad.__doc__ == "Computes the triad a*b+c." + t = triad.triad(12, 4, 15) assert triad.triad_called == 1 assert triad.triad_cache[(12, 4, 15)] == t