Completion for __getattr__ and __dir__ in module (PY-28017)

python/psi-api/src/com/jetbrains/python/PyNames.java

@@ -463,6 +463,12 @@ public class PyNames {
     .put("__fspath__", _only_self_descr)
     .build();
 
+  @NotNull
+  private static final ImmutableMap<String, BuiltinDescription> PY37_MODULE_BUILTIN_METHODS = ImmutableMap.<String, BuiltinDescription>builder()
+    .put("__getattr__", new BuiltinDescription("(name)"))
+    .put("__dir__", new BuiltinDescription("()"))
+    .build();
+
   public static ImmutableMap<String, BuiltinDescription> getBuiltinMethods(LanguageLevel level) {
     if (level.isAtLeast(LanguageLevel.PYTHON36)) {
       return PY36_BUILTIN_METHODS;
@@ -478,6 +484,15 @@ public class PyNames {
     }
   }
 
+  @NotNull
+  public static ImmutableMap<String, BuiltinDescription> getModuleBuiltinMethods(@NotNull LanguageLevel level) {
+    if (level.isAtLeast(LanguageLevel.PYTHON37)) {
+      return PY37_MODULE_BUILTIN_METHODS;
+    }
+
+    return ImmutableMap.of();
+  }
+
   // canonical names, not forced by interpreter
   public static final String CANONICAL_SELF = "self";
   public static final String CANONICAL_CLS = "cls";

python/src/com/jetbrains/extensions/python/CaptureExt.kt

@@ -15,4 +15,12 @@ import com.jetbrains.python.psi.PyFunction
 fun PsiElementPattern.Capture<*>.afterDefInMethod() =
   withLanguage(PythonLanguage.getInstance())
     .and(psiElement().inside(psiElement(PyFunction::class.java).inside(psiElement(PyClass::class.java))))
-                               .and(psiElement().afterLeaf("def"))
+    .and(psiElement().afterLeaf("def"))
+
+/**
+ * Place right after "def" in function or method. Useful to provide custom completions
+ */
+fun PsiElementPattern.Capture<*>.afterDefInFunction() =
+  withLanguage(PythonLanguage.getInstance())
+    .and(psiElement().inside(psiElement(PyFunction::class.java)))
+    .and(psiElement().afterLeaf("def"))

python/src/com/jetbrains/python/codeInsight/completion/PyCompletionUtils.kt

@@ -11,8 +11,12 @@ import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.codeInsight.lookup.TailTypeDecorator
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
+import com.jetbrains.python.psi.AccessDirection
 import com.jetbrains.python.psi.PyClass
 import com.jetbrains.python.psi.PyFunction
+import com.jetbrains.python.psi.PyTypedElement
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.PyType
 import com.jetbrains.python.psi.types.TypeEvalContext
 import icons.PythonIcons
 
@@ -35,6 +39,7 @@ fun autoInsertSingleItem(context: AutoCompletionContext) =
 
 
 fun CompletionParameters.getPyClass() = (ScopeUtil.getScopeOwner(position) as? PyFunction)?.containingClass
+fun CompletionParameters.getFile() = (ScopeUtil.getScopeOwner(position) as? PyFunction)?.containingFile
 
 
 fun CompletionParameters.getTypeEvalContext() = TypeEvalContext.codeCompletion(originalFile.project, originalFile)
@@ -43,19 +48,24 @@ fun CompletionParameters.getTypeEvalContext() = TypeEvalContext.codeCompletion(o
 /**
  * Add method completion to to result.
  * @param result destination
- * @param pyClass if provided will check if method does not exist already
+ * @param element if provided will check if method does not exist already
  * @param builderPostprocessor function to be used to tune lookup builder
  */
 fun addMethodToResult(result: CompletionResultSet,
-                      pyClass: PyClass?,
+                      element: PyTypedElement?,
                       typeEvalContext: TypeEvalContext,
                       methodName: String,
                       methodParentheses: String = "(self)",
                       builderPostprocessor: ((LookupElementBuilder) -> LookupElementBuilder)? = null) {
+  val type = if (element != null) typeEvalContext.getType(element) else null
+  if (type != null) {
+    val resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(typeEvalContext)
 
-  if (pyClass?.findMethodByName(methodName, false, typeEvalContext) != null) {
-    return
+    if (type.resolveMember(methodName, null, AccessDirection.READ, resolveContext)?.firstOrNull()?.element != null) {
+      return
+    }
   }
+
   val item = LookupElementBuilder.create(methodName + methodParentheses)
     .withIcon(PythonIcons.Python.Nodes.Cyan_dot)
   result.addElement(TailTypeDecorator.withTail(builderPostprocessor?.invoke(item) ?: item, TailType.CASE_COLON))

python/src/com/jetbrains/python/codeInsight/completion/PySpecialMethodNamesCompletionContributor.kt

@@ -4,26 +4,41 @@ package com.jetbrains.python.codeInsight.completion
 import com.intellij.codeInsight.completion.*
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.util.ProcessingContext
-import com.jetbrains.extensions.python.afterDefInMethod
+import com.jetbrains.extensions.python.afterDefInFunction
 import com.jetbrains.python.PyNames
 import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyTypedElement
 
 class PySpecialMethodNamesCompletionContributor : CompletionContributor() {
   override fun handleAutoCompletionPossibility(context: AutoCompletionContext) = autoInsertSingleItem(context)
 
   init {
-    extend(CompletionType.BASIC, PlatformPatterns.psiElement().afterDefInMethod(), MyCompletionProvider)
+    extend(CompletionType.BASIC, PlatformPatterns.psiElement().afterDefInFunction(), MyCompletionProvider)
   }
 
   private object MyCompletionProvider : CompletionProvider<CompletionParameters>() {
     override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext?, result: CompletionResultSet) {
-      val pyClass = parameters.getPyClass() ?: return
       val typeEvalContext = parameters.getTypeEvalContext()
 
-      PyNames.getBuiltinMethods(LanguageLevel.forElement(pyClass))
-        ?.forEach {
-          addMethodToResult(result, pyClass, typeEvalContext, it.key, it.value.signature) { it.withTypeText("predefined") }
+      val pyClass = parameters.getPyClass()
+      if (pyClass != null) {
+        PyNames.getBuiltinMethods(LanguageLevel.forElement(pyClass))
+          ?.forEach {
+            addMethodToResult(result, pyClass, typeEvalContext, it.key, it.value.signature) { it.withTypeText("predefined") }
+          }
+      }
+      else {
+        val file = parameters.getFile()
+        if (file != null) {
+          PyNames
+            .getModuleBuiltinMethods(LanguageLevel.forElement(file))
+            .forEach {
+              addMethodToResult(result, file as? PyTypedElement, typeEvalContext, it.key, it.value.signature) {
+                it.withTypeText("predefined")
+              }
+            }
         }
+      }
     }
   }
 }

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -1199,6 +1199,18 @@ public class PythonCompletionTest extends PyTestCase {
     doMultiFileTest();
   }
 
+  // PY-28017
+  public void testModuleGetAttrAndDir() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () -> {
+        final List<String> suggested = doTestByText("def __<caret>");
+        assertNotNull(suggested);
+        assertContainsElements(suggested, "__getattr__(name)", "__dir__()");
+      }
+    );
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion";