From 6b26265b32163d0911d3b61f038b4141bf3b7b4a Mon Sep 17 00:00:00 2001 From: Logan Hunt <39638017+dosisod@users.noreply.github.com> Date: Fri, 12 Jan 2024 18:53:32 -0800 Subject: [PATCH] Add `itemgetter()` support for FURB118 (#322): Closes #321. --- docs/checks.md | 22 ++++++ refurb/checks/common.py | 12 +++ refurb/checks/readability/use_operators.py | 90 ++++++++++++++++++---- refurb/gen.py | 3 +- refurb/main.py | 5 +- test/data/err_118.py | 11 +++ test/data/err_118.txt | 3 + test/data/stringify.py | 11 +++ test/data/stringify.txt | 9 +++ 9 files changed, 147 insertions(+), 19 deletions(-) diff --git a/docs/checks.md b/docs/checks.md index ef7178d..ad6aa67 100644 --- 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 index 95d2cc9..5bcc097 100644 --- 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 index 0c86ea5..2e625b1 100644 --- 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 index 5d216dd..25a98aa 100644 --- 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 index 67c569f..1de0361 100644 --- 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 index 30e2004..d6a3df5 100644 --- a/test/data/err_118.py +++ b/test/data/err_118.py @@ -34,6 +34,11 @@ 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 @@ -41,3 +46,9 @@ def f2(x): 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 index ab6fbf4..366d7a8 100644 --- 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 index 7a624da..a5297e4 100644 --- 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 index 378affa..ff5367b 100644 --- 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))`