rewrite lru_cache with maxsize=None and typed=False

Co-authored-by: Max Rozentsveyg <maxr@outlook.com>

pyupgrade/_plugins/lru_cache.py

@@ -2,6 +2,7 @@
 import functools
 from typing import Iterable
 from typing import List
+from typing import Optional
 from typing import Tuple
 
 from tokenize_rt import Offset
@@ -23,6 +24,34 @@ def _remove_call(i: int, tokens: List[Token]) -> None:
     del tokens[i:j + 1]
 
 
+def _is_literal_kwarg(
+        keyword: ast.keyword, name: str, value: Optional[bool],
+) -> bool:
+    return (
+        keyword.arg == name and
+        isinstance(keyword.value, ast.NameConstant) and
+        keyword.value.value is value
+    )
+
+
+def _eligible(keywords: List[ast.keyword]) -> bool:
+    if len(keywords) == 1:
+        return _is_literal_kwarg(keywords[0], 'maxsize', None)
+    elif len(keywords) == 2:
+        return (
+            (
+                _is_literal_kwarg(keywords[0], 'maxsize', None) and
+                _is_literal_kwarg(keywords[1], 'typed', False)
+            ) or
+            (
+                _is_literal_kwarg(keywords[1], 'maxsize', None) and
+                _is_literal_kwarg(keywords[0], 'typed', False)
+            )
+        )
+    else:
+        return False
+
+
 @register(ast.Call)
 def visit_Call(
         state: State,
@@ -48,10 +77,7 @@ def visit_Call(
             isinstance(node.func.value, ast.Name) and
             node.func.value.id == 'functools' and
             not node.args and
-            len(node.keywords) == 1 and
-            node.keywords[0].arg == 'maxsize' and
-            isinstance(node.keywords[0].value, ast.NameConstant) and
-            node.keywords[0].value.value is None
+            _eligible(node.keywords)
     ):
         func = functools.partial(
             find_and_replace_call, template='functools.cache',

tests/features/lru_cache_test.py

@@ -101,6 +101,22 @@ def test_fix_no_arg_decorators(s, expected):
             (3, 9),
             id='unrelated parameter',
         ),
+        pytest.param(
+            'import functools\n\n'
+            '@functools.lru_cache(maxsize=None, typed=True)\n'
+            'def foo():\n'
+            '    pass\n',
+            (3, 9),
+            id='typed=True',
+        ),
+        pytest.param(
+            'import functools\n\n'
+            '@functools.lru_cache(maxsize=None, typed=False, foo=False)\n'
+            'def foo():\n'
+            '    pass\n',
+            (3, 9),
+            id='invalid keyword',
+        ),
     ),
 )
 def test_fix_maxsize_none_decorators_noop(s, min_version):
@@ -121,6 +137,28 @@ def test_fix_maxsize_none_decorators_noop(s, min_version):
             '    pass\n',
             id='call with attr',
         ),
+        pytest.param(
+            'import functools\n\n'
+            '@functools.lru_cache(maxsize=None, typed=False)\n'
+            'def foo():\n'
+            '    pass\n',
+            'import functools\n\n'
+            '@functools.cache\n'
+            'def foo():\n'
+            '    pass\n',
+            id='call with attr, maxsize=None then typed=False',
+        ),
+        pytest.param(
+            'import functools\n\n'
+            '@functools.lru_cache(typed=False, maxsize=None)\n'
+            'def foo():\n'
+            '    pass\n',
+            'import functools\n\n'
+            '@functools.cache\n'
+            'def foo():\n'
+            '    pass\n',
+            id='call with attr, typed=False then maxsize=None',
+        ),
     ),
 )
 def test_fix_maxsize_none_decorators(s, expected):