diff --git a/docs/checks.md b/docs/checks.md index 9947271..4ac1d3f 100644 --- a/docs/checks.md +++ b/docs/checks.md @@ -2343,4 +2343,25 @@ Good: names = ["Bob", "Alice", "Charlie"] names.reverse() +``` + +## FURB188: `remove-prefix-or-suffix` + +Categories: `performance` `readability` `string` + +Don't explicitly check a string prefix/suffix if you're only going to +remove it, use `.removeprefix()` or `.removesuffix()` instead. + +Bad: + +```python +def strip_txt_extension(filename: str) -> str: + return filename[:-4] if filename.endswith(".txt") else filename +``` + +Good: + +```python +def strip_txt_extension(filename: str) -> str: + return filename.removesuffix(".txt") ``` \ No newline at end of file diff --git a/refurb/checks/common.py b/refurb/checks/common.py index 6f0f1f4..de3c311 100644 --- a/refurb/checks/common.py +++ b/refurb/checks/common.py @@ -9,12 +9,14 @@ CallExpr, ComparisonExpr, ComplexExpr, + ConditionalExpr, DictExpr, DictionaryComprehension, Expression, FloatExpr, ForStmt, GeneratorExpr, + IfStmt, IndexExpr, IntExpr, LambdaExpr, @@ -412,6 +414,12 @@ def _stringify(node: Node) -> str: case AssignmentStmt(lvalues=[lhs], rvalue=rhs): return f"{stringify(lhs)} = {stringify(rhs)}" + case IfStmt(expr=[expr], body=[Block(body=[body])], else_body=None): + return f"if {_stringify(expr)}: {_stringify(body)}" + + case ConditionalExpr(if_expr=if_true, cond=cond, else_expr=if_false): + return f"{_stringify(if_true)} if {_stringify(cond)} else {_stringify(if_false)}" + raise ValueError diff --git a/refurb/checks/string/remove_prefix_or_suffix.py b/refurb/checks/string/remove_prefix_or_suffix.py new file mode 100644 index 0000000..a3da9c7 --- /dev/null +++ b/refurb/checks/string/remove_prefix_or_suffix.py @@ -0,0 +1,197 @@ +from dataclasses import dataclass + +from mypy.nodes import ( + AssignmentStmt, + CallExpr, + ConditionalExpr, + Expression, + IfStmt, + IndexExpr, + IntExpr, + MemberExpr, + RefExpr, + SliceExpr, + StrExpr, + UnaryExpr, +) + +from refurb.checks.common import is_equivalent, stringify +from refurb.error import Error +from refurb.settings import Settings +from refurb.visitor.traverser import TraverserVisitor + + +@dataclass +class ErrorInfo(Error): + """ + Don't explicitly check a string prefix/suffix if you're only going to + remove it, use `.removeprefix()` or `.removesuffix()` instead. + + Bad: + + ``` + def strip_txt_extension(filename: str) -> str: + return filename[:-4] if filename.endswith(".txt") else filename + ``` + + Good: + + ``` + def strip_txt_extension(filename: str) -> str: + return filename.removesuffix(".txt") + ``` + """ + + name = "remove-prefix-or-suffix" + categories = ("performance", "readability", "string") + code = 188 + + +def does_expr_match_slice_amount(str_func: str, lhs: Expression, rhs: Expression) -> bool: + match str_func, lhs, rhs: + case ( + "startswith", + StrExpr(value=value), + SliceExpr(begin_index=IntExpr(value=str_len), end_index=None), + ) if len(value) == str_len: + return True + + case ( + "startswith", + value, + SliceExpr( + begin_index=CallExpr(callee=RefExpr(fullname="builtins.len"), args=[len_arg]), + end_index=None, + ), + ) if is_equivalent(value, len_arg): + return True + + case ( + "endswith", + StrExpr(value=value), + SliceExpr( + begin_index=None, + end_index=UnaryExpr(op="-", expr=IntExpr(value=str_len)), + ), + ) if len(value) == str_len: + return True + + case ( + "endswith", + value, + SliceExpr( + begin_index=None, + end_index=UnaryExpr( + op="-", + expr=CallExpr(callee=RefExpr(fullname="builtins.len"), args=[len_arg]), + ), + ), + ) if is_equivalent(value, len_arg): + return True + + return False + + +STR_FUNC_TO_REMOVE_FUNC = {"endswith": "removesuffix", "startswith": "removeprefix"} + + +ignored_nodes = set[int]() + + +class IgnoreElifNodes(TraverserVisitor): + def visit_if_stmt(self, o: IfStmt) -> None: + ignored_nodes.add(id(o)) + + if o.else_body: + if o.else_body.body: + else_body = o.else_body.body[0] + + if isinstance(else_body, IfStmt): + ignored_nodes.add(id(else_body)) + + self.accept(else_body) + + +def check(node: ConditionalExpr | IfStmt, errors: list[Error], settings: Settings) -> None: + if settings.get_python_version() < (3, 9): + return # pragma: no cover + + if id(node) in ignored_nodes: + return + + if isinstance(node, IfStmt): + if node.else_body: + IgnoreElifNodes().accept(node.else_body) + + return + + expr = node.expr[0] + body = node.body[0].body + + if len(body) != 1: + return + + match expr: + case CallExpr( + callee=MemberExpr( + expr=func_lhs, + name="endswith" | "startswith" as func_name, + ), + args=[func_arg], + ): + pass + + case _: + return + + match body[0]: + case AssignmentStmt( + lvalues=[lvalue], + rvalue=IndexExpr( + base=slice_lhs, + index=SliceExpr(stride=None) as slice_expr, + ), + ) if ( + is_equivalent(slice_lhs, func_lhs) + and is_equivalent(lvalue, slice_lhs) + and does_expr_match_slice_amount(func_name, func_arg, slice_expr) + ): + parts = [ + f"{stringify(slice_lhs)} = {stringify(slice_lhs)}.", + STR_FUNC_TO_REMOVE_FUNC[func_name], + f"({stringify(func_arg)})", + ] + + msg = f"Replace `{stringify(node)}` with `{''.join(parts)}`" + + errors.append(ErrorInfo.from_node(node, msg)) + + if isinstance(node, ConditionalExpr): + match node: + case ConditionalExpr( + if_expr=IndexExpr( + base=slice_lhs, + index=SliceExpr(stride=None) as slice_expr, + ), + cond=CallExpr( + callee=MemberExpr( + expr=func_lhs, + name="endswith" | "startswith" as func_name, + ), + args=[func_arg], + ), + else_expr=if_false, + ) if ( + is_equivalent(slice_lhs, func_lhs) + and is_equivalent(func_lhs, if_false) + and does_expr_match_slice_amount(func_name, func_arg, slice_expr) + ): + parts = [ + f"{stringify(slice_lhs)}.", + STR_FUNC_TO_REMOVE_FUNC[func_name], + f"({stringify(func_arg)})", + ] + + msg = f"Replace `{stringify(node)}` with `{''.join(parts)}`" + + errors.append(ErrorInfo.from_node(node, msg)) diff --git a/test/data/err_188.py b/test/data/err_188.py new file mode 100644 index 0000000..950780c --- /dev/null +++ b/test/data/err_188.py @@ -0,0 +1,130 @@ +# these should match + +def remove_extension_via_slice(filename: str) -> str: + if filename.endswith(".txt"): + filename = filename[:-4] + + return filename + + +def remove_extension_via_slice_len(filename: str, extension: str) -> str: + if filename.endswith(extension): + filename = filename[:-len(extension)] + + return filename + + +def remove_extension_via_ternary(filename: str) -> str: + return filename[:-4] if filename.endswith(".txt") else filename + + +def remove_extension_via_ternary_with_len(filename: str, extension: str) -> str: + return filename[:-len(extension)] if filename.endswith(extension) else filename + + +def remove_prefix(filename: str) -> str: + return filename[4:] if filename.startswith("abc-") else filename + + +def remove_prefix_via_len(filename: str, prefix: str) -> str: + return filename[len(prefix):] if filename.startswith(prefix) else filename + + +# these should not + +def remove_extension_with_mismatched_len(filename: str) -> str: + if filename.endswith(".txt"): + filename = filename[:3] + + return filename + + +def remove_extension_assign_to_different_var(filename: str) -> str: + if filename.endswith(".txt"): + other_var = filename[:-4] + + return filename + + +def remove_extension_with_multiple_stmts(filename: str) -> str: + if filename.endswith(".txt"): + print("do some work") + + filename = filename[:-4] + + if filename.endswith(".txt"): + filename = filename[:-4] + + print("do some work") + + return filename + + +def remove_extension_from_unrelated_var(filename: str) -> str: + xyz = "abc.txt" + + if filename.endswith(".txt"): + filename = xyz[:-4] + + return filename + + +def remove_extension_in_elif(filename: str) -> str: + if filename: + pass + + elif filename.endswith(".txt"): + filename = filename[:-4] + + return filename + + +def remove_extension_in_multiple_elif(filename: str) -> str: + if filename: + pass + + elif filename: + pass + + elif filename.endswith(".txt"): + filename = filename[:-4] + + return filename + + +def remove_extension_in_if_with_else(filename: str) -> str: + if filename.endswith(".txt"): + filename = filename[:-4] + + else: + pass + + return filename + + +def remove_extension_ternary_name_mismatch(filename: str): + xyz = "" + + _ = xyz[:-4] if filename.endswith(".txt") else filename + _ = filename[:-4] if xyz.endswith(".txt") else filename + _ = filename[:-4] if filename.endswith(".txt") else xyz + + +def remove_extension_slice_amount_mismatch(filename: str) -> None: + extension = ".txt" + + _ = filename[:-1] if filename.endswith(".txt") else filename + _ = filename[:-1] if filename.endswith(extension) else filename + _ = filename[:-len("")] if filename.endswith(extension) else filename + + +def remove_prefix_size_mismatch(filename: str) -> str: + return filename[3:] if filename.startswith("abc-") else filename + + +def remove_prefix_name_mismatch(filename: str) -> None: + xyz = "" + + _ = xyz[4:] if filename.startswith("abc-") else filename + _ = filename[4:] if xyz.startswith("abc-") else filename + _ = filename[4:] if filename.startswith("abc-") else xyz diff --git a/test/data/err_188.txt b/test/data/err_188.txt new file mode 100644 index 0000000..9eb7159 --- /dev/null +++ b/test/data/err_188.txt @@ -0,0 +1,6 @@ +test/data/err_188.py:4:5 [FURB188]: Replace `if filename.endswith(".txt"): filename = filename[:-4]` with `filename = filename.removesuffix(".txt")` +test/data/err_188.py:11:5 [FURB188]: Replace `if filename.endswith(extension): filename = filename[:-len(extension)]` with `filename = filename.removesuffix(extension)` +test/data/err_188.py:18:12 [FURB188]: Replace `filename[:-4] if filename.endswith(".txt") else filename` with `filename.removesuffix(".txt")` +test/data/err_188.py:22:12 [FURB188]: Replace `filename[:-len(extension)] if filename.endswith(extension) else filename` with `filename.removesuffix(extension)` +test/data/err_188.py:26:12 [FURB188]: Replace `filename[4:] if filename.startswith("abc-") else filename` with `filename.removeprefix("abc-")` +test/data/err_188.py:30:12 [FURB188]: Replace `filename[len(prefix):] if filename.startswith(prefix) else filename` with `filename.removeprefix(prefix)`