Issue 13227: Option to make the lru_cache() type specific (suggested by Andrew Koenig).

3 files changed, 36 insertions, 9 deletions

diff --git a/Lib/functools.py b/Lib/functools.py
index 038f284..1abb37a 100644
--- a/Lib/functools.py
+++ b/Lib/functools.py
@@ -121,12 +121,16 @@ except ImportError:
 
 _CacheInfo = namedtuple("CacheInfo", "hits misses maxsize currsize")
 
-def lru_cache(maxsize=100):
+def lru_cache(maxsize=100, typed=False):
     """Least-recently-used cache decorator.
 
     If *maxsize* is set to None, the LRU features are disabled and the cache
     can grow without bound.
 
+    If *typed* is True, arguments of different types will be cached separately.
+    For example, f(3.0) and f(3) will be treated as distinct calls with
+    distinct results.
+
     Arguments to the cached function must be hashable.
 
     View the cache statistics named tuple (hits, misses, maxsize, currsize) with
@@ -142,7 +146,7 @@ def lru_cache(maxsize=100):
     # to allow the implementation to change (including a possible C version).
 
     def decorating_function(user_function,
-                tuple=tuple, sorted=sorted, len=len, KeyError=KeyError):
+            *, tuple=tuple, sorted=sorted, map=map, len=len, type=type, KeyError=KeyError):
 
         hits = misses = 0
         kwd_mark = (object(),)          # separates positional and keyword args
@@ -156,7 +160,12 @@ def lru_cache(maxsize=100):
                 nonlocal hits, misses
                 key = args
                 if kwds:
-                    key += kwd_mark + tuple(sorted(kwds.items()))
+                    sorted_items = tuple(sorted(kwds.items()))
+                    key += kwd_mark + sorted_items
+                if typed:
+                    key += tuple(map(type, args))
+                    if kwds:
+                        key += tuple(type(v) for k, v in sorted_items)
                 try:
                     result = cache[key]
                     hits += 1
@@ -177,7 +186,12 @@ def lru_cache(maxsize=100):
                 nonlocal hits, misses
                 key = args
                 if kwds:
-                    key += kwd_mark + tuple(sorted(kwds.items()))
+                    sorted_items = tuple(sorted(kwds.items()))
+                    key += kwd_mark + sorted_items
+                if typed:
+                    key += tuple(map(type, args))
+                    if kwds:
+                        key += tuple(type(v) for k, v in sorted_items)
                 with lock:
                     try:
                         result = cache[key]
diff --git a/Lib/re.py b/Lib/re.py
index cdf5976..4b90b3f 100644
--- a/Lib/re.py
+++ b/Lib/re.py
@@ -207,7 +207,7 @@ def compile(pattern, flags=0):
 
 def purge():
     "Clear the regular expression caches"
-    _compile_typed.cache_clear()
+    _compile.cache_clear()
     _compile_repl.cache_clear()
 
 def template(pattern, flags=0):
@@ -253,11 +253,8 @@ def escape(pattern):
 
 _pattern_type = type(sre_compile.compile("", 0))
 
+@functools.lru_cache(maxsize=500, typed=True)
 def _compile(pattern, flags):
-    return _compile_typed(type(pattern), pattern, flags)
-
-@functools.lru_cache(maxsize=500)
-def _compile_typed(text_bytes_type, pattern, flags):
     # internal: compile pattern
     if isinstance(pattern, _pattern_type):
         if flags:
diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
index a31c92e..c4d9fe6 100644
--- a/Lib/test/test_functools.py
+++ b/Lib/test/test_functools.py
@@ -734,6 +734,22 @@ class TestLRU(unittest.TestCase):
             with self.assertRaises(IndexError):
                 func(15)
 
+    def test_lru_with_types(self):
+        for maxsize in (None, 100):
+            @functools.lru_cache(maxsize=maxsize, typed=True)
+            def square(x):
+                return x * x
+            self.assertEqual(square(3), 9)
+            self.assertEqual(type(square(3)), type(9))
+            self.assertEqual(square(3.0), 9.0)
+            self.assertEqual(type(square(3.0)), type(9.0))
+            self.assertEqual(square(x=3), 9)
+            self.assertEqual(type(square(x=3)), type(9))
+            self.assertEqual(square(x=3.0), 9.0)
+            self.assertEqual(type(square(x=3.0)), type(9.0))
+            self.assertEqual(square.cache_info().hits, 4)
+            self.assertEqual(square.cache_info().misses, 4)
+
 def test_main(verbose=None):
     test_classes = (
         TestPartial,