From 91f27d8587eb3428d2c369e8a7d9c4b5453dd080 Mon Sep 17 00:00:00 2001
From: Ilia Zakoulov <ilia.zakoulov@jetbrains.com>
Date: Tue, 27 Aug 2024 14:51:54 +0200
Subject: [PATCH] PY-76629: Suppress PyProtectedMemberInspection if a member is
 defined in .pyi

Protected member should not be highlighted as a warning if it resolves to .pyi file.
We assume that everything in .pyi file is a public API.

GitOrigin-RevId: c8275f3e48e3cd69b1676de9b78606f28ea224c8
---
 .../src/com/jetbrains/python/pyi/PyiInspectionExtension.kt  | 6 ++++++
 .../accessProtectedProperty/AccessProtectedProperty.py      | 4 ++++
 .../pyi/inspections/accessProtectedProperty/test.py         | 4 ++++
 .../pyi/inspections/accessProtectedProperty/test.pyi        | 3 +++
 .../com/jetbrains/python/pyi/PyiInspectionsTest.java        | 4 ++++
 5 files changed, 21 insertions(+)
 create mode 100644 python/testData/pyi/inspections/accessProtectedProperty/AccessProtectedProperty.py
 create mode 100644 python/testData/pyi/inspections/accessProtectedProperty/test.py
 create mode 100644 python/testData/pyi/inspections/accessProtectedProperty/test.pyi

diff --git a/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiInspectionExtension.kt b/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiInspectionExtension.kt
index a9f95bc04a71..c08aa50d2a5d 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiInspectionExtension.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiInspectionExtension.kt
@@ -20,6 +20,7 @@ import com.jetbrains.python.inspections.PyInspectionExtension
 import com.jetbrains.python.psi.PyFromImportStatement
 import com.jetbrains.python.psi.PyImportElement
 import com.jetbrains.python.psi.PyImportStatement
+import com.jetbrains.python.psi.PyReferenceExpression
 import com.jetbrains.python.psi.types.TypeEvalContext
 
 class PyiInspectionExtension : PyInspectionExtension() {
@@ -35,4 +36,9 @@ class PyiInspectionExtension : PyInspectionExtension() {
     }
     return elements.isEmpty() || elements.any { it.asName != null }
   }
+
+  override fun ignoreProtectedSymbol(expression: PyReferenceExpression, context: TypeEvalContext): Boolean {
+    val referencedElement = expression.reference.resolve()
+    return referencedElement?.containingFile is PyiFile
+  }
 }
diff --git a/python/testData/pyi/inspections/accessProtectedProperty/AccessProtectedProperty.py b/python/testData/pyi/inspections/accessProtectedProperty/AccessProtectedProperty.py
new file mode 100644
index 000000000000..5e6adae356e4
--- /dev/null
+++ b/python/testData/pyi/inspections/accessProtectedProperty/AccessProtectedProperty.py
@@ -0,0 +1,4 @@
+from test import MyClass
+
+public_api = MyClass._public_api
+not_public_api = <weak_warning descr="Access to a protected member _not_public_api of a class">MyClass._not_public_api</weak_warning>
diff --git a/python/testData/pyi/inspections/accessProtectedProperty/test.py b/python/testData/pyi/inspections/accessProtectedProperty/test.py
new file mode 100644
index 000000000000..3c3bfaa5fa4f
--- /dev/null
+++ b/python/testData/pyi/inspections/accessProtectedProperty/test.py
@@ -0,0 +1,4 @@
+class MyClass:
+    def __init__(self):
+        self._public_api = "abacaba"
+        self._not_public_api = "qwerty"
diff --git a/python/testData/pyi/inspections/accessProtectedProperty/test.pyi b/python/testData/pyi/inspections/accessProtectedProperty/test.pyi
new file mode 100644
index 000000000000..978799c2dab2
--- /dev/null
+++ b/python/testData/pyi/inspections/accessProtectedProperty/test.pyi
@@ -0,0 +1,3 @@
+class MyClass:
+    def __init__(self) -> None: ...
+    _public_api: str
diff --git a/python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java b/python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java
index c84cc7bfe24c..5c6f9e909e0a 100644
--- a/python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java
+++ b/python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java
@@ -16,6 +16,7 @@
 package com.jetbrains.python.pyi;
 
 import com.intellij.codeInspection.LocalInspectionTool;
+import com.intellij.idea.TestFor;
 import com.intellij.openapi.Disposable;
 import com.intellij.openapi.util.Disposer;
 import com.intellij.psi.PsiDocumentManager;
@@ -149,4 +150,7 @@ public class PyiInspectionsTest extends PyTestCase {
   public void testMissedSuperInitCall() {
     doPyiTest(PyMissingConstructorInspection.class);
   }
+
+  @TestFor(issues = "PY-16477")
+  public void testAccessProtectedProperty() { doPyTest(PyProtectedMemberInspection.class); }
 }