Make a switch to original element when pyi-one has been found in order to determine if its deprecated or not (PY-38101)

GitOrigin-RevId: baf95e1db4db898d1e18051972f69799c95fa3b0

python/src/com/jetbrains/python/inspections/PyDeprecationInspection.java

@@ -24,6 +24,7 @@ import com.intellij.psi.PsiElementVisitor;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.PyKnownDecoratorUtil.KnownDecorator;
+import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -59,7 +60,7 @@ public class PyDeprecationInspection extends PyInspection {
         final PyExpression exceptClass = exceptPart.getExceptClass();
         if (exceptClass != null && "ImportError".equals(exceptClass.getText())) return;
       }
-      final PsiElement resolveResult = node.getReference(getResolveContext()).resolve();
+      final PsiElement resolveResult = resolve(node);
       final PyFromImportStatement importStatement = PsiTreeUtil.getParentOfType(node, PyFromImportStatement.class);
       if (importStatement != null) {
         final PsiElement element = importStatement.resolveImportSource();
@@ -67,7 +68,7 @@ public class PyDeprecationInspection extends PyInspection {
       }
       String deprecationMessage = null;
       if (resolveResult instanceof PyFunction) {
-        deprecationMessage = ((PyFunction) resolveResult).getDeprecationMessage();
+        deprecationMessage = ((PyFunction)resolveResult).getDeprecationMessage();
       }
       else if (resolveResult instanceof PyFile) {
         deprecationMessage = ((PyFile)resolveResult).getDeprecationMessage();
@@ -114,5 +115,11 @@ public class PyDeprecationInspection extends PyInspection {
         }
       }
     }
+
+    @Nullable
+    private PyElement resolve(@NotNull PyReferenceExpression node) {
+      final PyElement resolve = PyUtil.as(node.getReference(getResolveContext()).resolve(), PyElement.class);
+      return resolve == null ? null : PyiUtil.getOriginalElementOrLeaveAsIs(resolve, PyElement.class);
+    }
   }
 }

python/testData/deprecation/deprecatedElementInPyi/a.py

@@ -0,0 +1,3 @@
+import <warning descr="the deprecated module is deprecated; use a non-deprecated module instead">module</warning>
+
+<warning descr="the deprecated module is deprecated; use a non-deprecated module instead">module</warning>.<warning descr="this is deprecated">bar</warning>()

python/testData/deprecation/deprecatedElementInPyi/module.py

@@ -0,0 +1,7 @@
+import warnings
+warnings.warn("the deprecated module is deprecated; use a non-deprecated module instead",
+                DeprecationWarning, 2)
+
+def bar():
+    import warnings
+    warnings.warn("this is deprecated", DeprecationWarning, 2)

python/testData/deprecation/deprecatedElementInPyi/module.pyi

@@ -0,0 +1,2 @@
+def bar():
+    pass

python/testSrc/com/jetbrains/python/PyDeprecationTest.java

@@ -97,4 +97,12 @@ public class PyDeprecationTest extends PyTestCase {
     myFixture.configureByFile("a.py");
     myFixture.checkHighlighting(true, false, false);
   }
+
+  // PY-38101
+  public void testDeprecatedElementInPyi() {
+    myFixture.enableInspections(PyDeprecationInspection.class);
+    myFixture.copyDirectoryToProject("deprecation/deprecatedElementInPyi", "");
+    myFixture.configureByFile("a.py");
+    myFixture.checkHighlighting(true, false, false);
+  }
 }