[flake8-pyi] Fix inconsistent handling of forward references for __new__, __enter__, __aenter__ in PYI034 (#22798)

Author: denyszhak

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI034.py

@@ -385,3 +385,10 @@ def __new__(cls, /, name: str, bases: tuple[object, ...], attrs: dict[str, objec
 class MetaclassInWhichSelfCannotBeUsed8(django.db.models.base.ModelBase):
     def __new__(cls, name: builtins.str, bases: tuple, attributes: dict, /, **kw) -> MetaclassInWhichSelfCannotBeUsed8:
         ...
+
+
+class UsesStringizedForwardReferences:
+    def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+    def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+    async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+    def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034

crates/ruff_linter/src/rules/flake8_pyi/rules/non_self_return_type.rs

@@ -182,7 +182,7 @@ pub(crate) fn non_self_return_type(
 
     if is_async {
         if name == "__aenter__"
-            && is_name(returns, &class_def.name)
+            && is_name_or_stringized_name(returns, &class_def.name, checker)
             && !is_final(&class_def.decorator_list, semantic)
         {
             add_diagnostic(checker, stmt, returns, class_def, name);
@@ -198,7 +198,7 @@ pub(crate) fn non_self_return_type(
         return;
     }
 
-    if is_name(returns, &class_def.name) {
+    if is_name_or_stringized_name(returns, &class_def.name, checker) {
         if matches!(name, "__enter__" | "__new__") && !is_final(&class_def.decorator_list, semantic)
         {
             add_diagnostic(checker, stmt, returns, class_def, name);
@@ -327,6 +327,11 @@ fn is_name(expr: &ast::Expr, name: &str) -> bool {
     id.as_str() == name
 }
 
+/// Return `true` if the given expression resolves to the given name,
+fn is_name_or_stringized_name(expr: &ast::Expr, name: &str, checker: &Checker) -> bool {
+    checker.match_maybe_stringized_annotation(expr, |expr| is_name(expr, name))
+}
+
 /// Return `true` if the given expression resolves to `typing.Self`.
 fn is_self(expr: &ast::Expr, checker: &Checker) -> bool {
     checker.match_maybe_stringized_annotation(expr, |expr| {

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI034_PYI034.py.snap

@@ -472,3 +472,77 @@ help: Use `Self` as return type
 362 | 
 363 | # Test cases based on issue #20781 - metaclasses that triggers IsMetaclass::Maybe
 note: This is an unsafe fix and may change runtime behavior
+
+PYI034 [*] `__new__` methods in classes like `UsesStringizedForwardReferences` usually return `self` at runtime
+   --> PYI034.py:391:9
+    |
+390 | class UsesStringizedForwardReferences:
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+    |         ^^^^^^^
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+    |
+help: Use `Self` as return type
+388 | 
+389 | 
+390 | class UsesStringizedForwardReferences:
+    -     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+391 +     def __new__(cls) -> typing.Self: ...       # PYI034
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+note: This is an unsafe fix and may change runtime behavior
+
+PYI034 [*] `__enter__` methods in classes like `UsesStringizedForwardReferences` usually return `self` at runtime
+   --> PYI034.py:392:9
+    |
+390 | class UsesStringizedForwardReferences:
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+    |         ^^^^^^^^^
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+    |
+help: Use `Self` as return type
+389 | 
+390 | class UsesStringizedForwardReferences:
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+    -     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+392 +     def __enter__(self) -> typing.Self: ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+note: This is an unsafe fix and may change runtime behavior
+
+PYI034 [*] `__aenter__` methods in classes like `UsesStringizedForwardReferences` usually return `self` at runtime
+   --> PYI034.py:393:15
+    |
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+    |               ^^^^^^^^^^
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+    |
+help: Use `Self` as return type
+390 | class UsesStringizedForwardReferences:
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+    -     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+393 +     async def __aenter__(self) -> typing.Self: ...  # PYI034
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+note: This is an unsafe fix and may change runtime behavior
+
+PYI034 [*] `__iadd__` methods in classes like `UsesStringizedForwardReferences` usually return `self` at runtime
+   --> PYI034.py:394:9
+    |
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+394 |     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+    |         ^^^^^^^^
+    |
+help: Use `Self` as return type
+391 |     def __new__(cls) -> "UsesStringizedForwardReferences": ...       # PYI034
+392 |     def __enter__(self) -> "UsesStringizedForwardReferences": ...    # PYI034
+393 |     async def __aenter__(self) -> "UsesStringizedForwardReferences": ...  # PYI034
+    -     def __iadd__(self, other) -> "UsesStringizedForwardReferences": ...  # PYI034
+394 +     def __iadd__(self, other) -> typing.Self: ...  # PYI034
+note: This is an unsafe fix and may change runtime behavior