From a38c05bf13a9050aeb44d16e0d9870aab7f7c742 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 27 May 2024 21:44:24 -0400
Subject: [PATCH] Avoid recommending context manager in `__enter__`
 implementations (#11575)

## Summary

Closes https://github.com/astral-sh/ruff/issues/11567.
---
 .../test/fixtures/flake8_simplify/SIM115.py    | 12 ++++++++++++
 .../rules/open_file_with_context_handler.rs    | 18 +++++++++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM115.py b/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM115.py
index 6a6925c45a603..48dd0e8671ef1 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM115.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM115.py
@@ -46,3 +46,15 @@
 # OK (quick one-liner to clear file contents)
 open("filename", "w").close()
 pathlib.Path("filename").open("w").close()
+
+
+# OK (custom context manager)
+class MyFile:
+    def __init__(self, filename: str):
+        self.filename = filename
+
+    def __enter__(self):
+        self.file = open(self.filename)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.file.close()
diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/open_file_with_context_handler.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/open_file_with_context_handler.rs
index 3ae0034fcff81..fe578e2781905 100644
--- a/crates/ruff_linter/src/rules/flake8_simplify/rules/open_file_with_context_handler.rs
+++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/open_file_with_context_handler.rs
@@ -2,7 +2,7 @@ use ruff_python_ast::{self as ast, Expr, Stmt};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::{ScopeKind, SemanticModel};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -114,24 +114,27 @@ fn match_exit_stack(semantic: &SemanticModel) -> bool {
 
 /// Return `true` if `func` is the builtin `open` or `pathlib.Path(...).open`.
 fn is_open(semantic: &SemanticModel, func: &Expr) -> bool {
-    // open(...)
+    // Ex) `open(...)`
     if semantic.match_builtin_expr(func, "open") {
         return true;
     }
 
-    // pathlib.Path(...).open()
+    // Ex) `pathlib.Path(...).open()`
     let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func else {
         return false;
     };
+
     if attr != "open" {
         return false;
     }
+
     let Expr::Call(ast::ExprCall {
         func: value_func, ..
     }) = &**value
     else {
         return false;
     };
+
     semantic
         .resolve_qualified_name(value_func)
         .is_some_and(|qualified_name| matches!(qualified_name.segments(), ["pathlib", "Path"]))
@@ -189,6 +192,15 @@ pub(crate) fn open_file_with_context_handler(checker: &mut Checker, func: &Expr)
         return;
     }
 
+    // Ex) `def __enter__(self): ...`
+    if let ScopeKind::Function(ast::StmtFunctionDef { name, .. }) =
+        &checker.semantic().current_scope().kind
+    {
+        if name == "__enter__" {
+            return;
+        }
+    }
+
     checker
         .diagnostics
         .push(Diagnostic::new(OpenFileWithContextHandler, func.range()));