From 5a86e1f87ee0b0894e641898e9b4aa0259474f83 Mon Sep 17 00:00:00 2001
From: cevans87 <c.d.evans87@gmail.com>
Date: Wed, 27 Feb 2019 15:45:26 -0800
Subject: [PATCH] memoize decorator now works with classes

---
 README.md                      |  9 +++++++++
 atools/decorator_mixin.py      |  5 ++++-
 atools/memoize_decorator.py    | 32 +++++++++++++++++++++++++++++++-
 setup.py                       |  2 +-
 test/test_memoize_decorator.py | 30 ++++++++++++++++++++++++++++++
 5 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 14eba0d..a32181e 100644
--- a/README.md
+++ b/README.md
@@ -32,6 +32,15 @@ Python 3.7+ async-enabled decorators and tools including
             other nine calls in this example wait for the result.
             await asyncio.gather(*[foo(1) for _ in range(10)])
 
+        - Classes may be memoized.
+            @memoize
+            Class Foo:
+                def init(self, _): ...
+
+            Foo(1)  # Instance is actually created.
+            Foo(1)  # Instance not created. Previously-cached instance returned.
+            Foo(2)  # Instance is actually created.
+
         - Calls to foo(1), foo(bar=1), and foo(1, baz='baz') are equivalent and only cached once
             @memoize
             def foo(bar, baz='baz'): ...
diff --git a/atools/decorator_mixin.py b/atools/decorator_mixin.py
index 8982dd4..8e600a1 100644
--- a/atools/decorator_mixin.py
+++ b/atools/decorator_mixin.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from functools import wraps
+import inspect
 from types import FunctionType
 from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, Type, Union
 
@@ -27,7 +28,9 @@ def __call__(cls, _decoratee: Optional[Decoratee] = None, **kwargs) -> Decorated
 
         decorator = super().__call__(_decoratee, **kwargs)
 
-        if not isinstance(_decoratee, FunctionType):
+        if inspect.isclass(decorator):
+            decorated = decorator
+        elif inspect.isclass(_decoratee):
             # _decoratee is a class. Our decorator should have already done its work.
             decorated = _decoratee
         else:
diff --git a/atools/memoize_decorator.py b/atools/memoize_decorator.py
index dd5c890..8564b7d 100644
--- a/atools/memoize_decorator.py
+++ b/atools/memoize_decorator.py
@@ -282,6 +282,15 @@ async def foo(bar) -> Any: ...
             other nine calls in this example wait for the result.
             await asyncio.gather(*[foo(1) for _ in range(10)])
 
+        - Classes may be memoized.
+            @memoize
+            Class Foo:
+                def init(self, _): ...
+
+            Foo(1)  # Instance is actually created.
+            Foo(1)  # Instance not created. Previously-cached instance returned.
+            Foo(2)  # Instance is actually created.
+
         - Calls to foo(1), foo(bar=1), and foo(1, baz='baz') are equivalent and only cached once
             @memoize
             def foo(bar, baz='baz'): ...
@@ -346,6 +355,27 @@ def foo(self): -> Any: ...
             a.bar(1)  # Foo.bar(a, 1) is actually called cached and again.
     """
 
+    def __new__(
+            cls,
+            fn: Fn,
+            *,
+            size: Optional[int] = None,
+            duration: Optional[Union[int, timedelta]] = None,
+    ):
+        if not inspect.isclass(fn):
+            return super().__new__(cls)
+
+        class WrappedMeta(type(fn)):
+            # noinspection PyMethodParameters
+            @memoize
+            def __call__(cls, *args, **kwargs):
+                return super().__call__(*args, **kwargs)
+
+        class Wrapped(fn, metaclass=WrappedMeta):
+            pass
+
+        return type(fn.__name__, (Wrapped,), {})
+
     def __init__(
             self,
             fn: Fn,
@@ -355,7 +385,7 @@ def __init__(
     ) -> None:
         if inspect.iscoroutinefunction(fn):
             self._memo = _MemoizeAsync(fn, size=size, duration=duration)
-        else:
+        elif not inspect.isclass(fn):
             self._memo = _MemoizeSync(fn, size=size, duration=duration)
 
     def __call__(self, *args, **kwargs) -> Any:
diff --git a/setup.py b/setup.py
index b53ef58..e562b8e 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='atools',
-    version='0.4.1',
+    version='0.5.0',
     packages=['', 'atools'],
     python_requires='>=3.7',
     url='https://github.com/cevans87/atools',
diff --git a/test/test_memoize_decorator.py b/test/test_memoize_decorator.py
index 36261b7..954c198 100644
--- a/test/test_memoize_decorator.py
+++ b/test/test_memoize_decorator.py
@@ -355,6 +355,36 @@ def test_async_no_event_loop_does_not_raise(self) -> None:
         async def foo() -> None:
             ...
 
+    def test_memoizes_class(self) -> None:
+        body = MagicMock()
+
+        class Bar:
+            ...
+
+        @memoize
+        class Foo(Bar):
+            def __init__(self, foo):
+                body(foo)
+
+        self.assertIs(Foo(0), Foo(0))
+        body.assert_called_once_with(0)
+        self.assertIsNot(Foo(0), Foo(1))
+
+    def test_memoizes_class_with_metaclass(self) -> None:
+        body = MagicMock()
+
+        class FooMeta(type):
+            pass
+
+        @memoize
+        class Foo(metaclass=FooMeta):
+            def __init__(self, foo):
+                body(foo)
+
+        self.assertIs(Foo(0), Foo(0))
+        body.assert_called_once_with(0)
+        self.assertIsNot(Foo(0), Foo(1))
+
 
 if __name__ == '__main__':
     unittest.main(verbosity=2)