Add itemgetter() support for FURB118 (#322):

Closes #321.
dosisod · Jan 13, 2024 · 6b26265 · 6b26265
1 parent 389e0ff
commit 6b26265
Show file tree

Hide file tree

Showing 9 changed files with 147 additions and 19 deletions.
diff --git a/docs/checks.md b/docs/checks.md
@@ -434,6 +434,28 @@ nums = [1, 2, 3]
 print(reduce(add, nums))  # 6
 ```
 
+In addition, the `operator.itemgetter()` function can be used to get one or
+more items from an object, removing the need to create a lambda just to
+extract values from an object:
+
+Bad:
+
+```python
+row = (1, "Some text", True)
+
+transform = lambda x: (x[2], x[0])
+```
+
+Good:
+
+```python
+from operator import itemgetter
+
+row = (1, "Some text", True)
+
+transform = itemgetter(2, 0)
+```
+
 ## FURB119: `use-fstring-format`
 
 Categories: `builtin` `fstring`

diff --git a/refurb/checks/common.py b/refurb/checks/common.py
@@ -388,3 +388,15 @@ def _stringify(node: Node) -> str:
             return f"[{inner}]"
 
     raise ValueError
+
+
+def slice_expr_to_slice_call(expr: SliceExpr) -> str:
+    args = [
+        stringify(expr.begin_index) if expr.begin_index else "None",
+        stringify(expr.end_index) if expr.end_index else "None",
+    ]
+
+    if expr.stride:
+        args.append(stringify(expr.stride))
+
+    return f"slice({', '.join(args)})"
diff --git a/refurb/checks/readability/use_operators.py b/refurb/checks/readability/use_operators.py
@@ -5,14 +5,17 @@
     Block,
     ComparisonExpr,
     FuncItem,
+    IndexExpr,
     LambdaExpr,
     NameExpr,
     OpExpr,
     ReturnStmt,
+    SliceExpr,
+    TupleExpr,
     UnaryExpr,
 )
 
-from refurb.checks.common import _stringify
+from refurb.checks.common import _stringify, slice_expr_to_slice_call, stringify
 from refurb.error import Error
 
 
@@ -42,6 +45,28 @@ class ErrorInfo(Error):
 
     print(reduce(add, nums))  # 6
     ```
+
+    In addition, the `operator.itemgetter()` function can be used to get one or
+    more items from an object, removing the need to create a lambda just to
+    extract values from an object:
+
+    Bad:
+
+    ```
+    row = (1, "Some text", True)
+
+    transform = lambda x: (x[2], x[0])
+    ```
+
+    Good:
+
+    ```
+    from operator import itemgetter
+
+    row = (1, "Some text", True)
+
+    transform = itemgetter(2, 0)
+    ```
     """
 
     name = "use-operator"
@@ -123,24 +148,57 @@ def check(node: FuncItem, errors: list[Error]) -> None:
         case FuncItem(
             arg_names=[name],
             arg_kinds=[ArgKind.ARG_POS],
-            body=Block(
-                body=[
-                    ReturnStmt(
-                        expr=UnaryExpr(
-                            op=op,
-                            expr=NameExpr(name=expr_name),
+            body=Block(body=[ReturnStmt(expr=expr)]),
+        ):
+            match expr:
+                case UnaryExpr(
+                    op=op,
+                    expr=NameExpr(name=expr_name),
+                ) if name == expr_name and (func_name := UNARY_OPERATORS.get(op)):
+                    errors.append(
+                        ErrorInfo.from_node(
+                            node,
+                            f"Replace {func_type} with `operator.{func_name}`",
                         )
                     )
-                ]
-            ),
-        ) if name == expr_name:
-            if func_name := UNARY_OPERATORS.get(op):
-                errors.append(
-                    ErrorInfo.from_node(
-                        node,
-                        f"Replace {func_type} with `operator.{func_name}`",
+
+                case IndexExpr(
+                    base=NameExpr(name=item_name),
+                    index=index,
+                ) if item_name == name:
+                    index_expr = (
+                        slice_expr_to_slice_call(index)
+                        if isinstance(index, SliceExpr)
+                        else stringify(index)
                     )
-                )
+
+                    msg = f"Replace {func_type} with `operator.itemgetter({index_expr})`"
+
+                    errors.append(ErrorInfo.from_node(node, msg))
+
+                case TupleExpr(items=items) if items:
+                    tuple_args: list[str] = []
+
+                    for item in items:
+                        match item:
+                            case IndexExpr(
+                                base=NameExpr(name=item_name),
+                                index=index,
+                            ) if item_name == name:
+                                tuple_args.append(
+                                    slice_expr_to_slice_call(index)
+                                    if isinstance(index, SliceExpr)
+                                    else stringify(index)
+                                )
+
+                            case _:
+                                return
+
+                    inner = ", ".join(tuple_args)
+
+                    msg = f"Replace {func_type} with `operator.itemgetter({inner})`"
+
+                    errors.append(ErrorInfo.from_node(node, msg))
 
 
 def get_function_type(node: FuncItem) -> str:

diff --git a/refurb/gen.py b/refurb/gen.py
@@ -2,6 +2,7 @@
 import sys
 from collections import defaultdict
 from contextlib import suppress
+from operator import itemgetter
 from pathlib import Path
 from subprocess import PIPE, run
 
@@ -127,7 +128,7 @@ def build_imports(names: list[str]) -> str:
 
     return "\n".join(
         f"from {module} import {', '.join(names)}"
-        for module, names in sorted(modules.items(), key=lambda x: x[0])
+        for module, names in sorted(modules.items(), key=itemgetter(0))
     )
 
 

diff --git a/refurb/main.py b/refurb/main.py
@@ -6,6 +6,7 @@
 from functools import cache, partial
 from importlib import metadata
 from io import StringIO
+from operator import itemgetter
 from pathlib import Path
 from tempfile import mkstemp
 
@@ -333,12 +334,12 @@ def output_timing_stats(
     data = {
         "mypy_total_time_spent_in_ms": int(mypy_total_time_spent * 1_000),
         "mypy_time_spent_parsing_modules_in_ms": dict(
-            sorted(mypy_stats.items(), key=lambda x: x[1], reverse=True)
+            sorted(mypy_stats.items(), key=itemgetter(1), reverse=True)
         ),
         "refurb_time_spent_checking_file_in_ms": dict(
             sorted(
                 refurb_timing_stats_in_ms.items(),
-                key=lambda x: x[1],
+                key=itemgetter(1),
                 reverse=True,
             )
         ),

diff --git a/test/data/err_118.py b/test/data/err_118.py
@@ -34,10 +34,21 @@ def f(x, y):
 def f2(x):
     return - x
 
+lambda x: x[0]
+
+lambda x: (x[0], x[1], x[2])
+lambda x: (x[1:], x[2])
+
 
 # these will not
 
 lambda x, y: print(x + y)
 lambda x, *y: x + y
 lambda x, y: y + x
 lambda x, y: 1 + 2
+lambda x: (1, x[1], x[2])
+lambda x: (x.y, x[1], x[2])
+lambda x, y: (x[0], y[0])
+lambda x, y: (x[0], y[0])
+lambda x: ()
+lambda x: (*x[0], x[1])
diff --git a/test/data/err_118.txt b/test/data/err_118.txt
@@ -26,3 +26,6 @@ test/data/err_118.py:28:1 [FURB118]: Replace `lambda x: not x` with `operator.no
 test/data/err_118.py:29:1 [FURB118]: Replace `lambda x: +x` with `operator.pos`
 test/data/err_118.py:31:1 [FURB118]: Replace function with `operator.add`
 test/data/err_118.py:34:1 [FURB118]: Replace function with `operator.neg`
+test/data/err_118.py:37:1 [FURB118]: Replace `lambda x: x[0]` with `operator.itemgetter(0)`
+test/data/err_118.py:39:1 [FURB118]: Replace `lambda x: (x[0], x[1], x[2])` with `operator.itemgetter(0, 1, 2)`
+test/data/err_118.py:40:1 [FURB118]: Replace `lambda x: (x[1:], x[2])` with `operator.itemgetter(slice(1, None), 2)`
diff --git a/test/data/stringify.py b/test/data/stringify.py
@@ -7,3 +7,14 @@
 _ = list([""[1:2]])
 _ = list([""[1:2:-1]])
 _ = list([""[::-1]])
+
+# test slice exprs
+lambda x: x[1:]
+lambda x: x[1:2]
+lambda x: x[1:2:3]
+lambda x: x[1::3]
+lambda x: x[:2:3]
+lambda x: x[::3]
+lambda x: x[:2:]
+lambda x: x[::]
+lambda x: x[:]
diff --git a/test/data/stringify.txt b/test/data/stringify.txt
@@ -4,3 +4,12 @@ test/data/stringify.py:6:5 [FURB123]: Replace `list([""[1:]])` with `[""[1:]].co
 test/data/stringify.py:7:5 [FURB123]: Replace `list([""[1:2]])` with `[""[1:2]].copy()`
 test/data/stringify.py:8:5 [FURB123]: Replace `list([""[1:2:-1]])` with `[""[1:2:-1]].copy()`
 test/data/stringify.py:9:5 [FURB123]: Replace `list([""[::-1]])` with `[""[::-1]].copy()`
+test/data/stringify.py:12:1 [FURB118]: Replace `lambda x: x[1:]` with `operator.itemgetter(slice(1, None))`
+test/data/stringify.py:13:1 [FURB118]: Replace `lambda x: x[1:2]` with `operator.itemgetter(slice(1, 2))`
+test/data/stringify.py:14:1 [FURB118]: Replace `lambda x: x[1:2:3]` with `operator.itemgetter(slice(1, 2, 3))`
+test/data/stringify.py:15:1 [FURB118]: Replace `lambda x: x[1::3]` with `operator.itemgetter(slice(1, None, 3))`
+test/data/stringify.py:16:1 [FURB118]: Replace `lambda x: x[:2:3]` with `operator.itemgetter(slice(None, 2, 3))`
+test/data/stringify.py:17:1 [FURB118]: Replace `lambda x: x[::3]` with `operator.itemgetter(slice(None, None, 3))`
+test/data/stringify.py:18:1 [FURB118]: Replace `lambda x: x[:2]` with `operator.itemgetter(slice(None, 2))`
+test/data/stringify.py:19:1 [FURB118]: Replace `lambda x: x[:]` with `operator.itemgetter(slice(None, None))`
+test/data/stringify.py:20:1 [FURB118]: Replace `lambda x: x[:]` with `operator.itemgetter(slice(None, None))`