Add remove-prefix-or-suffix check

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")
 ```

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
 
 

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))

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

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)`