fixed PY-7857 Doctest: missing completion and highlighting for python keywords

python/src/com/jetbrains/python/codeInsight/completion/PyKeywordCompletionContributor.java

@@ -12,18 +12,21 @@ import com.intellij.patterns.StandardPatterns;
 import com.intellij.psi.*;
 import com.intellij.psi.filters.ElementFilter;
 import com.intellij.psi.filters.position.FilterPattern;
+import com.intellij.psi.impl.source.tree.injected.InjectedFileViewProvider;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ProcessingContext;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.codeInsight.UnindentingInsertHandler;
+import com.jetbrains.python.documentation.doctest.PyDocstringFile;
 import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import static com.intellij.patterns.PlatformPatterns.psiElement;
+import static com.intellij.patterns.PlatformPatterns.psiFile;
 import static com.intellij.patterns.StandardPatterns.or;
 
 /**
@@ -140,7 +143,10 @@ public class PyKeywordCompletionContributor extends CompletionContributor {
       if (p instanceof PsiComment) return false; // just in case
       int point = p.getTextOffset();
       PsiDocumentManager docMgr = PsiDocumentManager.getInstance(p.getProject());
-      Document doc = docMgr.getDocument(p.getContainingFile().getOriginalFile());
+      final PsiFile file = p.getContainingFile().getOriginalFile();
+      Document doc = docMgr.getDocument(file);
+      String indentCharacters = file.getViewProvider() instanceof InjectedFileViewProvider? " \t>": " \t";
+
       if (doc != null) {
         CharSequence chs = doc.getCharsSequence();
         char c;
@@ -150,7 +156,7 @@ public class PyKeywordCompletionContributor extends CompletionContributor {
           c = chs.charAt(point);
           if (c == '\n') return true;
         }
-        while (c == ' ' || c == '\t');
+        while (indentCharacters.indexOf(c) >= 0);
       }
       return false;
     }
@@ -213,7 +219,8 @@ public class PyKeywordCompletionContributor extends CompletionContributor {
     psiElement().inside(PsiComment.class);
 
   private static final PsiElementPattern.Capture<PsiElement> IN_STRING_LITERAL =
-    psiElement().inside(PyStringLiteralExpression.class);
+    psiElement().inside(PyStringLiteralExpression.class).andNot(
+      psiElement().inFile(psiFile(PyDocstringFile.class)));
 
   private static final ElementPattern<PsiElement> IN_FUNCTION_HEADER =
     or(psiElement().inside(PyFunction.class).andNot(psiElement().inside(false, psiElement(PyStatementList.class), psiElement(PyFunction.class))),

python/src/com/jetbrains/python/documentation/doctest/PyDocstringFile.java

@@ -0,0 +1,38 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.injection.InjectedLanguageManager;
+import com.intellij.openapi.fileTypes.FileType;
+import com.intellij.psi.FileViewProvider;
+import com.intellij.psi.PsiLanguageInjectionHost;
+import com.jetbrains.python.psi.LanguageLevel;
+import com.jetbrains.python.psi.impl.PyFileImpl;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User: ktisha
+ */
+public class PyDocstringFile extends PyFileImpl {
+
+  public PyDocstringFile(FileViewProvider viewProvider) {
+    super(viewProvider, PyDocstringLanguageDialect.getInstance());
+  }
+
+  @NotNull
+  @Override
+  public FileType getFileType() {
+    return PyDocstringFileType.INSTANCE;
+  }
+
+  @Override
+  public String toString() {
+    return "DocstringFile:" + getName();
+  }
+
+  @Override
+  public LanguageLevel getLanguageLevel() {
+    final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(getProject());
+    final PsiLanguageInjectionHost host = languageManager.getInjectionHost(this);
+    if (host != null) return LanguageLevel.forElement(host.getContainingFile());
+    return super.getLanguageLevel();
+  }
+}

python/src/com/jetbrains/python/documentation/doctest/PyDocstringParserDefinition.java

@@ -1,18 +1,13 @@
 package com.jetbrains.python.documentation.doctest;
 
 import com.intellij.lang.PsiParser;
-import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.lexer.Lexer;
-import com.intellij.openapi.fileTypes.FileType;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.FileViewProvider;
 import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.tree.IFileElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.jetbrains.python.PythonParserDefinition;
-import com.jetbrains.python.psi.LanguageLevel;
-import com.jetbrains.python.psi.impl.PyFileImpl;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -47,25 +42,6 @@ public class PyDocstringParserDefinition extends PythonParserDefinition {
 
   @Override
   public PsiFile createFile(FileViewProvider viewProvider) {
-    return new PyFileImpl(viewProvider, PyDocstringLanguageDialect.getInstance()) {
-      @NotNull
-      @Override
-      public FileType getFileType() {
-        return PyDocstringFileType.INSTANCE;
-      }
-
-      @Override
-      public String toString() {
-        return "DocstringFile:" + getName();
-      }
-
-      @Override
-      public LanguageLevel getLanguageLevel() {
-        final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(getProject());
-        final PsiLanguageInjectionHost host = languageManager.getInjectionHost(this);
-        if (host != null) return LanguageLevel.forElement(host.getContainingFile());
-        return super.getLanguageLevel();
-      }
-    };
+    return new PyDocstringFile(viewProvider);
   }
 }

python/testData/completion/doctest/forInDoctest.py

@@ -0,0 +1,6 @@
+__author__ = 'ktisha'
+
+def bar():
+	"""
+	>>> fo<caret>
+	"""

python/testData/completion/doctest/functionInDoctest.py

@@ -0,0 +1,6 @@
+__author__ = 'ktisha'
+
+def foo():
+	"""
+	>>> fo<caret>
+	"""

python/testData/completion/doctest/importInDoctest.py

@@ -0,0 +1,6 @@
+__author__ = 'ktisha'
+from functionInDoctest import foo
+def bar():
+	"""
+	>>> fo<caret>
+	"""

python/testSrc/com/jetbrains/python/PythonDoctestCompletionTest.java

@@ -0,0 +1,37 @@
+package com.jetbrains.python;
+
+import com.intellij.codeInsight.lookup.LookupElement;
+import com.jetbrains.python.fixtures.PyTestCase;
+
+/**
+ * User : ktisha
+ */
+public class PythonDoctestCompletionTest extends PyTestCase {
+
+  private void doDoctestTest(String expected) {
+    final String testName = "completion/doctest/" + getTestName(true);
+    myFixture.configureByFile(testName + ".py");
+    final LookupElement[] elements = myFixture.completeBasic();
+    if (elements != null) {
+      for (LookupElement lookup : elements) {
+        System.out.println(lookup.getLookupString());
+        if (lookup.getLookupString().equals(expected))
+          return;
+      }
+    }
+    fail();
+  }
+
+  public void testForInDoctest() {
+    doDoctestTest("for");
+  }
+
+  public void testImportInDoctest() {
+    doDoctestTest("foo");
+  }
+
+  public void testFunctionInDoctest() {
+    doDoctestTest("foo");
+  }
+
+}