initial doctests support

2026-09-13 21:55:01 +07:00 · 2012-10-16 14:33:26 +04:00
parent 0700051973
commit a898ec3db7
22 changed files with 634 additions and 3 deletions
@@ -683,6 +683,15 @@

    <problemFileHighlightFilter implementation="com.jetbrains.python.buildout.config.BuildoutCfgProblemFileHighlightFilter"/>

+
+    <!-- PyDocstring -->
+    <languageInjector implementation="com.jetbrains.python.documentation.doctest.PyDocstringLanguageInjector"/>
+    <lang.parserDefinition language="PyDocstring" implementationClass="com.jetbrains.python.documentation.doctest.PyDocstringParserDefinition"/>
+
+
+
+    <lang.syntaxHighlighterFactory key="PyDocstring" implementationClass="com.jetbrains.python.documentation.doctest.PyDocstringSyntaxHighlighterFactory"/>
+
    <!-- Mako files -->
    <multiLangCommenter implementation="com.jetbrains.python.templateLanguages.TemplatesCommentProvider"/>
    <highlightErrorFilter implementation="com.jetbrains.mako.inspection.MakoErrorFilter"/>
@@ -838,6 +847,10 @@
    <dialectsTokenSetContributor implementation="com.jetbrains.mako.MakoTokenSetContributor"/>
    <visitorFilter language="Mako" implementationClass="com.jetbrains.mako.highlighting.MakoVisitorFilter"/>

+    <!-- PyDocstring -->
+    <dialectsTokenSetContributor implementation="com.jetbrains.python.documentation.doctest.PyDocstringTokenSetContributor"/>
+    <visitorFilter language="PyDocstring" implementationClass="com.jetbrains.python.documentation.doctest.PyDocstringVisitorFilter"/>
+
    <!-- Console -->
    <visitorFilter language="Python" implementationClass="com.jetbrains.python.console.ConsoleVisitorFilter"/>

@@ -0,0 +1,71 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.google.common.collect.Lists;
+import com.intellij.codeInsight.completion.CompletionUtil;
+import com.intellij.lang.injection.InjectedLanguageManager;
+import com.intellij.psi.*;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.psi.PyQualifiedExpression;
+import com.jetbrains.python.psi.impl.ResolveResultList;
+import com.jetbrains.python.psi.impl.references.PyReferenceImpl;
+import com.jetbrains.python.psi.resolve.*;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.ArrayList;
+
+/**
+ * User : ktisha
+ */
+public class PyDocReference extends PyReferenceImpl {
+  public PyDocReference(PyQualifiedExpression element, @NotNull PyResolveContext context) {
+    super(element, context);
+  }
+
+
+  @NotNull
+  @Override
+  public ResolveResult[] multiResolve(boolean incompleteCode) {
+    ResolveResult[] results = super.multiResolve(incompleteCode);
+    if (results.length == 0) {
+      final ResolveResultList ret = new ResolveResultList();
+      PsiFile file = myElement.getContainingFile();
+      final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myElement.getProject());
+      final PsiLanguageInjectionHost host = languageManager.getInjectionHost(myElement);
+      if (host != null) file = host.getContainingFile();
+
+      final String referencedName = myElement.getReferencedName();
+      if (referencedName == null) return ResolveResult.EMPTY_ARRAY;
+      ResolveProcessor processor = new ResolveProcessor(referencedName);
+
+      PyResolveUtil.scopeCrawlUp(processor, (ScopeOwner)file, referencedName, file);
+      PsiElement uexpr = processor.getResult();
+      if (uexpr != null) ret.add(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, uexpr));
+      if (ret.size() > 0) {
+        return ret.toArray(new RatedResolveResult[ret.size()]);
+      }
+      return new ResolveResult[] { new PsiElementResolveResult(myElement) };
+    }
+    return results;
+  }
+
+  @NotNull
+  public Object[] getVariants() {
+    final ArrayList<Object> ret = Lists.newArrayList(super.getVariants());
+    PsiFile file = myElement.getContainingFile();
+    final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myElement.getProject());
+    final PsiLanguageInjectionHost host = languageManager.getInjectionHost(myElement);
+    if (host != null) file = host.getContainingFile();
+
+    final PsiElement originalElement = CompletionUtil.getOriginalElement(myElement);
+    final PyQualifiedExpression element = originalElement instanceof PyQualifiedExpression ?
+                                          (PyQualifiedExpression)originalElement : myElement;
+
+    // include our own names
+    final CompletionVariantsProcessor processor = new CompletionVariantsProcessor(element);
+    PyResolveUtil.scopeCrawlUp(processor, (ScopeOwner)file, null, null);
+
+    ret.addAll(processor.getResultList());
+
+    return ret.toArray();
+  }
+}
@@ -0,0 +1,30 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiPolyVariantReference;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.impl.PyReferenceExpressionImpl;
+import com.jetbrains.python.psi.impl.references.PyQualifiedReference;
+import com.jetbrains.python.psi.resolve.PyResolveContext;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ *
+ * User : ktisha
+ */
+public class PyDocReferenceExpression extends PyReferenceExpressionImpl {
+
+  public PyDocReferenceExpression(ASTNode astNode) {
+    super(astNode);
+  }
+
+  @NotNull
+  public PsiPolyVariantReference getReference(PyResolveContext context) {
+    final PyExpression qualifier = getQualifier();
+    if (qualifier != null) {
+      return new PyQualifiedReference(this, context);
+    }
+    return new PyDocReference(this, context);
+  }
+}
+
@@ -0,0 +1,17 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.PsiBuilder;
+import com.jetbrains.python.parsing.ParsingContext;
+import com.jetbrains.python.parsing.PyParser;
+import com.jetbrains.python.parsing.StatementParsing;
+import com.jetbrains.python.psi.LanguageLevel;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstingParser extends PyParser {
+  @Override
+  protected ParsingContext createParsingContext(PsiBuilder builder, LanguageLevel languageLevel, StatementParsing.FUTURE futureFlag) {
+    return new PyDocstringParsingContext(builder, languageLevel, futureFlag);
+  }
+}
@@ -0,0 +1,18 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.Language;
+import com.jetbrains.python.psi.PyFileElementType;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringFileElementType extends PyFileElementType {
+  public PyDocstringFileElementType(Language language) {
+    super(language);
+  }
+
+  @Override
+  public String getExternalId() {
+    return "PyDocstring.FILE";
+  }
+}
@@ -0,0 +1,34 @@
+
+package com.jetbrains.python.documentation.doctest;
+
+import com.jetbrains.python.PythonFileType;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringFileType extends PythonFileType {
+  public static PythonFileType INSTANCE = new PyDocstringFileType();
+
+  protected PyDocstringFileType() {
+    super(new PyDocstringLanguageDialect());
+  }
+
+  @NotNull
+  @Override
+  public String getName() {
+    return "PyDocstring";
+  }
+
+  @NotNull
+  @Override
+  public String getDescription() {
+    return "python docstring";
+  }
+
+  @NotNull
+  @Override
+  public String getDefaultExtension() {
+    return "docstring";
+  }
+}
@@ -0,0 +1,17 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.Language;
+import com.jetbrains.python.PythonLanguage;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringLanguageDialect extends Language {
+  public static PyDocstringLanguageDialect getInstance() {
+    return (PyDocstringLanguageDialect)PyDocstringFileType.INSTANCE.getLanguage();
+  }
+
+  protected PyDocstringLanguageDialect() {
+    super(PythonLanguage.getInstance(), "PyDocstring");
+  }
+}
@@ -0,0 +1,70 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.InjectedLanguagePlaces;
+import com.intellij.psi.LanguageInjector;
+import com.intellij.psi.PsiLanguageInjectionHost;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.PyDocStringOwner;
+import com.jetbrains.python.psi.PyStringLiteralExpression;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+
+/**
+ * User: ktisha
+ */
+public class PyDocstringLanguageInjector implements LanguageInjector {
+  @Override
+  public void getLanguagesToInject(@NotNull final PsiLanguageInjectionHost host, @NotNull final InjectedLanguagePlaces injectionPlacesRegistrar) {
+    if (host instanceof PyStringLiteralExpression) {
+      final PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(host, PyDocStringOwner.class);
+      if (docStringOwner != null) {
+        if (docStringOwner.getDocStringExpression() == host) {
+
+          int start = 0;
+          int end = host.getTextLength() - 1;
+          final String text = host.getText();
+          final List<String> strings = StringUtil.split(text, "\n", false);
+
+          boolean gotExample = false;
+
+          int currentPosition = 0;
+          for (String string : strings) {
+            final String trimmedString = string.trim();
+            if (!trimmedString.startsWith(">>>") && !trimmedString.startsWith("...") && gotExample && start < end) {
+              gotExample = false;
+              injectionPlacesRegistrar.addPlace(PyDocstringLanguageDialect.getInstance(), TextRange.create(start, end-1),  null, null);
+            }
+            if (trimmedString.startsWith(">>>")) {
+              if (!gotExample)
+                start = currentPosition;
+
+              gotExample = true;
+              end = getEndOffset(currentPosition, string);
+            }
+            else if (trimmedString.startsWith("...") && gotExample) {
+              end = getEndOffset(currentPosition, string);
+            }
+            currentPosition = currentPosition + string.length();
+          }
+          if (gotExample && start < end)
+            injectionPlacesRegistrar.addPlace(PyDocstringLanguageDialect.getInstance(), TextRange.create(start, end-1),  null, null);
+        }
+      }
+    }
+  }
+
+  private int getEndOffset(int start, String s) {
+    int end;
+    int length = s.length();
+    if (s.trim().endsWith("\"\"\"") || s.trim().endsWith("'''"))
+      length = length - 3;
+    else if (s.trim().endsWith("\"") || s.trim().endsWith("'"))
+      length = length - 1;
+
+    end = start + length;
+    return end;
+  }
+}
@@ -0,0 +1,32 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.lexer.PythonIndentingLexer;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringLexer extends PythonIndentingLexer {
+  @Override
+  public void advance() {
+    if (super.getTokenType() == PyTokenTypes.DOT) {
+      advanceBase();
+      if (super.getTokenType() == PyTokenTypes.DOT) {
+        advanceBase();
+        if (super.getTokenType() == PyTokenTypes.DOT) {
+          advanceBase();
+        }
+      }
+    }
+    else if (super.getTokenType() == PyTokenTypes.GTGT) {
+      super.advance();
+      if (super.getTokenType() == PyTokenTypes.GT) {
+        super.advance();
+      }
+    }
+    else {
+      super.advance();
+    }
+  }
+
+}
@@ -0,0 +1,55 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.PsiParser;
+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.tree.IFileElementType;
+import com.intellij.psi.tree.TokenSet;
+import com.jetbrains.python.PythonParserDefinition;
+import com.jetbrains.python.psi.impl.PyFileImpl;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringParserDefinition extends PythonParserDefinition {
+  public static final IFileElementType PYTHON_DOCSTRING_FILE = new PyDocstringFileElementType(PyDocstringLanguageDialect
+                                                                                                .getInstance());
+
+  @NotNull
+  public Lexer createLexer(Project project) {
+    return new PyDocstringLexer();
+  }
+
+  @NotNull
+  @Override
+  public PsiParser createParser(Project project) {
+    return new PyDocstingParser();
+  }
+
+
+  @NotNull
+  @Override
+  public TokenSet getWhitespaceTokens() {
+    return TokenSet.orSet(super.getWhitespaceTokens(), TokenSet.create(PyDocstringTokenTypes.DOTS));
+  }
+
+  @Override
+  public IFileElementType getFileNodeType() {
+    return PYTHON_DOCSTRING_FILE;
+  }
+
+  @Override
+  public PsiFile createFile(FileViewProvider viewProvider) {
+    return new PyFileImpl(viewProvider, PyDocstringLanguageDialect.getInstance()) {
+      @NotNull
+      @Override
+      public FileType getFileType() {
+        return PyDocstringFileType.INSTANCE;
+      }
+    };
+  }
+}
@@ -0,0 +1,82 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.lang.PsiBuilder;
+import com.intellij.psi.tree.IElementType;
+import com.intellij.util.text.CharArrayUtil;
+import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.parsing.ExpressionParsing;
+import com.jetbrains.python.parsing.ParsingContext;
+import com.jetbrains.python.parsing.ParsingScope;
+import com.jetbrains.python.parsing.StatementParsing;
+import com.jetbrains.python.psi.LanguageLevel;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringParsingContext extends ParsingContext {
+  private final StatementParsing stmtParser;
+  private final ExpressionParsing exprParser;
+
+  public PyDocstringParsingContext(final PsiBuilder builder,
+                                   LanguageLevel languageLevel,
+                                   StatementParsing.FUTURE futureFlag) {
+    super(builder, languageLevel, futureFlag);
+    stmtParser = new PyDocstringStatementParsing(this, futureFlag);
+    exprParser = new PyDocstringExpressionParsing(this);
+  }
+
+  @Override
+  public ExpressionParsing getExpressionParser() {
+    return exprParser;
+  }
+
+  @Override
+  public StatementParsing getStatementParser() {
+    return stmtParser;
+  }
+
+  private static class PyDocstringExpressionParsing extends ExpressionParsing {
+    public PyDocstringExpressionParsing(ParsingContext context) {
+      super(context);
+    }
+
+    @Override
+    protected IElementType getReferenceType() {
+      return PyDocstringTokenTypes.DOC_REFERENCE;
+    }
+  }
+
+  private static class PyDocstringStatementParsing extends StatementParsing {
+
+    protected IElementType getReferenceType() {
+      return PyDocstringTokenTypes.DOC_REFERENCE;
+    }
+
+    protected PyDocstringStatementParsing(ParsingContext context,
+                                          @Nullable FUTURE futureFlag) {
+      super(context, futureFlag);
+    }
+    @Override
+    public void parseStatement(ParsingScope scope) {
+      IElementType type = myBuilder.getTokenType();
+      while (!myBuilder.eof() && type != PyDocstringTokenTypes.WELCOME) {
+        myBuilder.advanceLexer();
+        type = myBuilder.getTokenType();
+      }
+      if (type == PyDocstringTokenTypes.WELCOME) { // >>> case
+        myBuilder.advanceLexer();
+        super.parseStatement(scope);
+      }
+    }
+
+    @Override
+    public IElementType filter(IElementType source, int start, int end, CharSequence text) {
+      if (source == PyTokenTypes.DOT && CharArrayUtil.regionMatches(text, start, end, "..."))
+        return PyDocstringTokenTypes.DOTS;
+      if (source == PyTokenTypes.GTGT && CharArrayUtil.regionMatches(text, start, end, ">>>"))
+        return PyDocstringTokenTypes.WELCOME;
+      return super.filter(source, start, end, text);
+    }
+  }
+}
@@ -0,0 +1,27 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.openapi.editor.colors.TextAttributesKey;
+import com.intellij.openapi.fileTypes.SyntaxHighlighter;
+import com.intellij.openapi.fileTypes.SyntaxHighlighterFactory;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.tree.IElementType;
+import com.jetbrains.python.highlighting.PyHighlighter;
+import com.jetbrains.python.psi.LanguageLevel;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringSyntaxHighlighterFactory extends SyntaxHighlighterFactory {
+  @NotNull
+  public SyntaxHighlighter getSyntaxHighlighter(final Project project, final VirtualFile virtualFile) {
+    return new PyHighlighter(LanguageLevel.getDefault()){
+      @NotNull
+      @Override
+      public TextAttributesKey[] getTokenHighlights(IElementType tokenType) {
+        return new TextAttributesKey[0];
+      }
+    };
+  }
+}
@@ -0,0 +1,14 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.psi.tree.TokenSet;
+import com.jetbrains.python.PythonTokenSetContributor;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringTokenSetContributor extends PythonTokenSetContributor {
+  @Override
+  public TokenSet getExpressionTokens() {
+    return TokenSet.orSet(super.getExpressionTokens(), TokenSet.create(PyDocstringTokenTypes.DOC_REFERENCE));
+  }
+}
@@ -0,0 +1,15 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.jetbrains.python.psi.PyElementType;
+
+/**
+ * User : ktisha
+ */
+public class PyDocstringTokenTypes {
+  public static final PyElementType DOC_REFERENCE = new PyElementType("DOC_REFERENCE", PyDocReferenceExpression.class);
+  public static final PyElementType WELCOME = new PyElementType("WELCOME");
+  public static final PyElementType DOTS = new PyElementType("DOTS");
+
+  private PyDocstringTokenTypes() {
+  }
+}
@@ -0,0 +1,34 @@
+package com.jetbrains.python.documentation.doctest;
+
+import com.intellij.psi.PsiFile;
+import com.jetbrains.python.inspections.*;
+import com.jetbrains.python.validation.DocStringAnnotator;
+import com.jetbrains.python.validation.HighlightingAnnotator;
+import com.jetbrains.python.validation.ParameterListAnnotator;
+import com.jetbrains.python.validation.ReturnAnnotator;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User : ktisha
+ *
+ * filter out some python inspections and annotations if we're in docstring substitution
+ */
+public class PyDocstringVisitorFilter implements PythonVisitorFilter {
+  @Override
+  public boolean isSupported(@NotNull final Class visitorClass, @NotNull final PsiFile file) {
+    //inspections
+    if (visitorClass == PyArgumentListInspection.class) {
+      return false;
+    }
+    if (visitorClass == PyDocstringInspection.class || visitorClass == PyStatementEffectInspection.class ||
+        visitorClass == PyUnboundLocalVariableInspection.class || visitorClass == PyUnnecessaryBackslashInspection.class ||
+        visitorClass == PyUnresolvedReferencesInspection.class) {
+      return false;
+    }
+    //annotators
+    if (visitorClass == DocStringAnnotator.class || visitorClass == ParameterListAnnotator.class || visitorClass == ReturnAnnotator.class ||
+        visitorClass == HighlightingAnnotator.class)
+      return false;
+    return true;
+  }
+}
@@ -355,7 +355,10 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt

    @Override
    public int getOffsetInHost(final int offsetInDecoded, @NotNull TextRange rangeInsideHost) {
-      return myHost.valueOffsetToTextOffset(offsetInDecoded);
+      int offset = offsetInDecoded + rangeInsideHost.getStartOffset();
+      if (offset < rangeInsideHost.getStartOffset()) offset = rangeInsideHost.getStartOffset();
+      if (offset > rangeInsideHost.getEndOffset()) offset = rangeInsideHost.getEndOffset();
+      return offset;
    }

    @Override
@@ -0,0 +1,8 @@
+class Test:
+  pass
+
+def foo():
+  """
+  >>> a = Tes<caret>
+
+  """
@@ -0,0 +1,8 @@
+class Test:
+  pass
+
+def foo():
+  """
+  >>> a = Test
+
+  """
@@ -0,0 +1,5 @@
+def foo():
+  """
+  >>> fo<caret>
+
+  """
@@ -0,0 +1,5 @@
+def foo():
+  """
+  >>> foo
+
+  """
@@ -0,0 +1,73 @@
+package com.jetbrains.python;
+
+import com.intellij.lexer.Lexer;
+import com.jetbrains.python.documentation.doctest.PyDocstringParserDefinition;
+import com.jetbrains.python.fixtures.PyTestCase;
+
+/**
+ * User: ktisha
+ */
+public class PyDocstringTest extends PyTestCase {
+
+  @Override
+  protected String getTestDataPath() {
+    return super.getTestDataPath() + "/doctests/";
+  }
+
+  public void testWelcome() {
+    doTestLexer("  >>> foo()", "Py:SPACE", "Py:INDENT", "Py:GTGT", "Py:SPACE", "Py:IDENTIFIER", "Py:LPAR", "Py:RPAR", "Py:STATEMENT_BREAK");
+  }
+
+  public void testDots() {
+    doTestLexer(">>> grouped == { 2:2,\n" +
+                "  ...              3:3}", "Py:GTGT","Py:SPACE", "Py:IDENTIFIER", "Py:SPACE", "Py:EQEQ", "Py:SPACE", "Py:LBRACE", "Py:SPACE", "Py:INTEGER_LITERAL", "Py:COLON", "Py:INTEGER_LITERAL", "Py:COMMA", "Py:LINE_BREAK", "Py:DOT", "Py:SPACE", "Py:SPACE", "Py:INTEGER_LITERAL", "Py:COLON", "Py:INTEGER_LITERAL", "Py:RBRACE", "Py:STATEMENT_BREAK");
+  }
+
+  public void testFunctionName() throws Throwable {
+    doCompletionTest();
+  }
+
+  public void testClassName() throws Throwable {
+    doCompletionTest();
+  }
+
+  public void doCompletionTest() throws Throwable {
+    String inputDataFileName = getInputDataFileName(getTestName(true));
+    String expectedResultFileName = getExpectedResultFileName(getTestName(true));
+    myFixture.testCompletion(inputDataFileName, expectedResultFileName);
+  }
+
+  // util methods
+  private static String getInputDataFileName(String testName) {
+    return testName + ".docstring";
+  }
+
+  private static String getExpectedResultFileName(String testName) {
+    return testName + ".expected.docstring";
+  }
+
+
+  private void doTestLexer(final String text, String... expectedTokens) {
+    Lexer lexer = new PyDocstringParserDefinition().createLexer(myFixture.getProject());
+    lexer.start(text);
+    int idx = 0;
+    while (lexer.getTokenType() != null) {
+      if (idx >= expectedTokens.length) {
+        StringBuilder remainingTokens = new StringBuilder("\"" + lexer.getTokenType().toString() + "\"");
+        lexer.advance();
+        while (lexer.getTokenType() != null) {
+          remainingTokens.append(",");
+          remainingTokens.append(" \"").append(lexer.getTokenType().toString()).append("\"");
+          lexer.advance();
+        }
+        fail("Too many tokens. Following tokens: " + remainingTokens.toString());
+      }
+      String tokenName = lexer.getTokenType().toString();
+      assertEquals("Token mismatch at position " + idx, expectedTokens[idx], tokenName);
+      idx++;
+      lexer.advance();
+    }
+
+    if (idx < expectedTokens.length) fail("Not enough tokens");
+  }
+}
@@ -26,8 +26,8 @@ public class PyStringLiteralTest extends PyTestCase {
    assertEquals("\n", builder.toString());

    assertEquals(1, escaper.getOffsetInHost(0, new TextRange(1, 5)));
-    assertEquals(3, escaper.getOffsetInHost(1, new TextRange(1, 5)));
-    assertEquals(6, escaper.getOffsetInHost(4, new TextRange(1, 5)));
+    assertEquals(2, escaper.getOffsetInHost(1, new TextRange(1, 5)));
+    assertEquals(5, escaper.getOffsetInHost(4, new TextRange(1, 5)));
  }

  private PyStringLiteralExpression createLiteralFromText(final String text) {