PY-10184 Impl hasattr support

GitOrigin-RevId: e5826372ddfe233656097da774d978c57ee54937

python/python-common-tests/com/jetbrains/python/PyCommonResolveTest.java

@@ -1465,4 +1465,18 @@ public abstract class PyCommonResolveTest extends PyCommonResolveTestCase {
     final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
     assertEmpty(file.findTopLevelAttribute("t").multiResolveAssignedValue(PyResolveContext.noImplicits().withTypeEvalContext(context)));
   }
+
+  // PY-10184
+  public void testHasattrResolveTrueIfBranch() {
+    PsiElement targetElement = resolve();
+    assertInstanceOf(targetElement, PyStringLiteralExpression.class);
+    assertEquals("ajjj", ((PyStringLiteralExpression)targetElement).getStringValue());
+  }
+
+  // PY-10184
+  public void testHasattrResolveConditionalExpression() {
+    PsiElement targetElement = resolve();
+    assertInstanceOf(targetElement, PyStringLiteralExpression.class);
+    assertEquals("fld", ((PyStringLiteralExpression)targetElement).getStringValue());
+  }
 }

python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java

@@ -1661,6 +1661,72 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
     assertTrue(indImport < ind__import__);
   }
 
+  // PY-10184
+  public void testHasattrSimpleAnd() {
+    String[] inList = {"foo", "bar"};
+    doTestHasattrContributor(inList, null);
+  }
+
+  // PY-10184
+  public void testHasattrAndWithNot() {
+    String[] inList = {"foo", "baz"};
+    String[] notInList = {"bar"};
+    doTestHasattrContributor(inList, notInList);
+  }
+
+  // PY-10184
+  public void testHasattrIfPyramidAndOr() {
+    String[] inList = {"foo", "qux", "corge", "bar", "baz", "quux", "quuz"};
+    doTestHasattrContributor(inList, null);
+  }
+
+  // PY-10184
+  public void testHasattrInElseBranch() {
+    String[] notInList = {"ajjj"};
+    doTestHasattrContributor(null, notInList);
+  }
+
+  // PY-10184
+  public void testHasattrInElseBranchAfterElif() {
+    String[] notInList = {"foo1", "foo2"};
+    doTestHasattrContributor(null, notInList);
+  }
+
+  // PY-10184
+  public void testHasattrInElifBranch() {
+    String[] inList = {"foo2"};
+    String[] notInList = {"foo1"};
+    doTestHasattrContributor(inList, notInList);
+  }
+
+  // PY-10184
+  public void testHasattrInRightPartOfAnd() {
+    String[] inList = {"foo", "bar"};
+    doTestHasattrContributor(inList, null);
+  }
+
+  // PY-10184
+  public void testHasattrInConditionalExpression() {
+    String[] inList = {"foo", "bar"};
+    doTestHasattrContributor(inList, null);
+  }
+
+  private void doTestHasattrContributor(String[] inList, String[] notInList) {
+    doTestHasattrContributor("hasattrCompletion/" + getTestName(true) + ".py", inList, notInList);
+  }
+
+  private void doTestHasattrContributor(String testFileName, String[] inList, String[] notInList) {
+    myFixture.configureByFile(testFileName);
+    myFixture.completeBasic();
+    List<String> suggested = myFixture.getLookupElementStrings();
+    if (inList != null) {
+      assertContainsElements(suggested, inList);
+    }
+    if (notInList != null) {
+      assertDoesntContain(suggested, notInList);
+    }
+  }
+
   private void assertNoVariantsInExtendedCompletion() {
     myFixture.copyDirectoryToProject(getTestName(true), "");
     myFixture.configureByFile("a.py");

python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java

@@ -23,7 +23,6 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Comparing;
-import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
 import com.intellij.psi.search.GlobalSearchScope;
@@ -44,6 +43,7 @@ import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyImportedModule;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.ResolveResultList;
+import com.jetbrains.python.psi.impl.references.hasattr.PyHasAttrHelper;
 import com.jetbrains.python.psi.resolve.ImplicitResolveResult;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
@@ -119,6 +119,9 @@ public class PyQualifiedReference extends PyReferenceImpl {
     if ("__doc__".equals(referencedName)) {
       addDocReference(ret, qualifier, qualifierType);
     }
+
+    PyHasAttrHelper.INSTANCE.addHasAttrResolveResults(myElement, referencedName, qualifier, ret);
+
     return ret;
   }
 
@@ -272,7 +275,6 @@ public class PyQualifiedReference extends PyReferenceImpl {
     ctx.put(PyType.CTX_NAMES, namesAlready);
     final Collection<Object> variants = new ArrayList<>();
     if (qualifierType != null) {
-      Collections.addAll(variants, getVariantFromHasAttr(qualifier));
       if (qualifierType instanceof PyStructuralType && ((PyStructuralType)qualifierType).isInferredFromUsages()) {
         final PyClassType guessedType = guessClassTypeByName();
         if (guessedType != null) {
@@ -298,7 +300,6 @@ public class PyQualifiedReference extends PyReferenceImpl {
           }
         }
         Collections.addAll(variants, qualifierType.getCompletionVariants(element.getName(), element, ctx));
-        return variants.toArray();
       }
       else {
         return qualifierType.getCompletionVariants(element.getName(), element, ctx);
@@ -312,24 +313,10 @@ public class PyQualifiedReference extends PyReferenceImpl {
       if (qualifier instanceof PyReferenceExpression) {
         Collections.addAll(variants, collectSeenMembers(qualifier.getText()));
       }
-      return variants.toArray();
     }
-  }
 
-  private Object[] getVariantFromHasAttr(PyExpression qualifier) {
-    Collection<Object> variants = new ArrayList<>();
-    PyIfStatement ifStatement = PsiTreeUtil.getParentOfType(myElement, PyIfStatement.class);
-    while (ifStatement != null) {
-      PyExpression condition = ifStatement.getIfPart().getCondition();
-      if (condition instanceof PyCallExpression && ((PyCallExpression)condition).isCalleeText(PyNames.HAS_ATTR)) {
-        PyCallExpression call = (PyCallExpression)condition;
-        if (call.getArguments().length > 1 && call.getArguments()[0].getText().equals(qualifier.getText())) {
-          PyStringLiteralExpression string = call.getArgument(1, PyStringLiteralExpression.class);
-          if (string != null && StringUtil.isJavaIdentifier(string.getStringValue())) variants.add(string.getStringValue());
-        }
-      }
-      ifStatement = PsiTreeUtil.getParentOfType(ifStatement, PyIfStatement.class);
-    }
+    PyHasAttrHelper.INSTANCE.addHasAttrCompletionResults(element, qualifier, namesAlready, variants);
+
     return variants.toArray();
   }
 

python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/hasattr/PyHasAttrHelper.kt

@@ -0,0 +1,87 @@
+package com.jetbrains.python.psi.impl.references.hasattr
+
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.PlatformIcons
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.ResolveResultList
+import com.jetbrains.python.psi.resolve.RatedResolveResult
+
+object PyHasAttrHelper {
+  fun addHasAttrResolveResults(psiElement: PsiElement, referencedName: String, qualifier: PyExpression, ret: ResolveResultList) {
+    val hasAttrVariants = getHasAttrVariantsFromContext(psiElement, qualifier)
+    for (variant in hasAttrVariants.keys) {
+      if (variant == referencedName) {
+        ret.add(RatedResolveResult(RatedResolveResult.RATE_NORMAL, hasAttrVariants[variant]))
+        return
+      }
+    }
+  }
+
+  fun addHasAttrCompletionResults(psiElement: PsiElement, qualifier: PyExpression,
+                                  namesAlready: MutableSet<String>, variants: MutableCollection<Any>) {
+    for (variant in variants) {
+      if (variant is LookupElement) {
+        namesAlready.add(variant.lookupString)
+      }
+      else {
+        namesAlready.add(variant.toString())
+      }
+    }
+
+    for (variant in getHasAttrVariantsFromContext(psiElement, qualifier).keys) {
+      if (!namesAlready.contains(variant)) {
+        variants.add(LookupElementBuilder.create(variant)
+                       .withTypeText(PyNames.HAS_ATTR)
+                       .withIcon(PlatformIcons.FIELD_ICON))
+        namesAlready.add(variant)
+      }
+    }
+  }
+
+  private fun getHasAttrVariantsFromContext(psiElement: PsiElement, qualifier: PyExpression): Map<String, PsiElement> {
+    val result = hashMapOf<String, PsiElement>()
+    result.putAll(getHasAttrVariantsFromAnd(psiElement, qualifier))
+    result.putAll(getHasAttrVariantsFromConditions(psiElement, qualifier))
+    return result
+  }
+
+  private fun getHasAttrVariantsFromAnd(psiElement: PsiElement, qualifier: PyExpression): Map<String, PsiElement> {
+    val result = hashMapOf<String, PsiElement>()
+    val binaryExpr = PsiTreeUtil.getParentOfType(psiElement, PyBinaryExpression::class.java) ?: return result
+    if (!binaryExpr.isOperator(PyNames.AND)) return result
+    if (!PsiTreeUtil.isAncestor(binaryExpr.rightExpression, psiElement, false)) return result
+    result.putAll(getHasAttrVisitorResultOn(binaryExpr.leftExpression, qualifier))
+    return result
+  }
+
+  private fun getHasAttrVariantsFromConditions(psiElement: PsiElement, qualifier: PyExpression): Map<String, PsiElement> {
+    val result = hashMapOf<String, PsiElement>()
+
+    var curParent = PsiTreeUtil.getParentOfType(psiElement, PyIfPart::class.java, PyConditionalExpression::class.java)
+    while (curParent != null) {
+      val condition = when {
+        curParent is PyIfPart && PsiTreeUtil.isAncestor(curParent.statementList, psiElement, true) -> curParent.condition
+        curParent is PyConditionalExpression && PsiTreeUtil.isAncestor(curParent.truePart, psiElement, false) -> curParent.condition
+        else -> null
+      }
+      if (condition != null) {
+        result.putAll(getHasAttrVisitorResultOn(condition, qualifier))
+      }
+      curParent = PsiTreeUtil.getParentOfType(curParent, PyIfPart::class.java, PyConditionalExpression::class.java)
+    }
+
+    return result
+  }
+
+  private fun getHasAttrVisitorResultOn(psiElement: PsiElement, qualifier: PyExpression): Map<String, PsiElement> {
+    if (qualifier !is PyReferenceExpression) return hashMapOf()
+    val resolvedQualifier = qualifier.reference.resolve() ?: return hashMapOf()
+    val pyHasAttrVisitor = PyHasAttrVisitor(resolvedQualifier)
+    psiElement.accept(pyHasAttrVisitor)
+    return pyHasAttrVisitor.result
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/hasattr/PyHasAttrVisitor.kt

@@ -0,0 +1,48 @@
+package com.jetbrains.python.psi.impl.references.hasattr
+
+import com.intellij.codeInsight.completion.CompletionUtilCoreImpl
+import com.intellij.openapi.util.text.StringUtil
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.PyTokenTypes
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+class PyHasAttrVisitor(private val resolvedQualifier: PsiElement) : PyRecursiveElementVisitor() {
+  val result = hashMapOf<String, PsiElement>()
+  private var myPositive: Boolean = true
+
+  override fun visitPyPrefixExpression(node: PyPrefixExpression) {
+    if (node.operator === PyTokenTypes.NOT_KEYWORD) {
+      myPositive = !myPositive
+      super.visitPyPrefixExpression(node)
+      myPositive = !myPositive
+    }
+    else {
+      super.visitPyPrefixExpression(node)
+    }
+  }
+
+  override fun visitPyBinaryExpression(node: PyBinaryExpression) {
+    if (!node.isOperator(PyNames.AND) && !node.isOperator(PyNames.OR)) return
+    super.visitPyBinaryExpression(node)
+  }
+
+  override fun visitPyCallExpression(node: PyCallExpression) {
+    if (!myPositive) return
+    if (!node.isCalleeText(PyNames.HAS_ATTR)) return
+    if (node.arguments.size != 2) return
+    val firstArg = node.getArgument(0, PyReferenceExpression::class.java) ?: return
+    val attrName = node.getArgument(1, PyStringLiteralExpression::class.java) ?: return
+
+    if (firstArg.reference.isReferenceTo(resolvedQualifier)) {
+      val variant = attrName.stringValue
+      if (StringUtil.isJavaIdentifier(variant)) {
+        if (!result.containsKey(variant)) {
+          result[variant] = attrName
+        }
+      }
+    }
+  }
+}

python/testData/completion/hasattrCompletion/hasattrAndWithNot.py

@@ -0,0 +1,3 @@
+def foo(x):
+    if hasattr(x, 'foo') and not hasattr(x, 'bar') and not not hasattr(x, 'baz'):
+        print(x.<caret>)

python/testData/completion/hasattrCompletion/hasattrIfPyramidAndOr.py

@@ -0,0 +1,5 @@
+def foo(x):
+    if hasattr(x, 'foo'):
+      if hasattr(x, 'bar') or hasattr(x, 'baz'):
+        if hasattr(x, 'qux') and (hasattr(x, 'quux') or hasattr(x, 'quuz')) and hasattr(x, 'corge'):
+          print(x.<caret>)

python/testData/completion/hasattrCompletion/hasattrInConditionalExpression.py

@@ -0,0 +1,2 @@
+def foo(x):
+  some_var = (x.<caret> if hasattr(x, "foo") else 42) if hasattr(x, "bar") else 42

python/testData/completion/hasattrCompletion/hasattrInElifBranch.py

@@ -0,0 +1,5 @@
+def foo(x: int):
+    if hasattr(x, 'foo1'):
+        print(x.foo1)
+    elif hasattr(x, 'foo2'):
+        print(x.<caret>)

python/testData/completion/hasattrCompletion/hasattrInElseBranch.py

@@ -0,0 +1,5 @@
+def foo(x: int):
+    if hasattr(x, 'ajjj'):
+        print(x.ajjj)
+    else:
+        print(x.<caret>

python/testData/completion/hasattrCompletion/hasattrInElseBranchAfterElif.py

@@ -0,0 +1,7 @@
+def foo(x: int):
+    if hasattr(x, 'foo1'):
+        print(x.foo1)
+    elif hasattr(x, 'foo2'):
+        print(x.foo)
+    else:
+        print(x.<caret>)

python/testData/completion/hasattrCompletion/hasattrInRightPartOfAnd.py

@@ -0,0 +1,3 @@
+class Test(object):
+    def is_test(self):
+        return hasattr(self, 'foo') and hasattr(self, 'bar') and self.<caret>

python/testData/completion/hasattrCompletion/hasattrSimpleAnd.py

@@ -0,0 +1,3 @@
+def foo(x):
+    if hasattr(x, 'foo') and hasattr(x, 'bar'):
+        print(x.<caret>)

python/testData/resolve/HasattrResolveConditionalExpression.py

@@ -0,0 +1,2 @@
+def foo(x):
+    return x.f<ref>ld if hasattr(x, "fld") else 42

python/testData/resolve/HasattrResolveTrueIfBranch.py

@@ -0,0 +1,3 @@
+def foo(x: int):
+    if hasattr(x, 'ajjj'):
+        print(x.a<ref>jjj)