Type assertions for assertIsInstance (PY-4383)

python/src/com/jetbrains/python/PyNames.java

@@ -49,6 +49,7 @@ public class PyNames {
   public static final String DEBUG = "__debug__";
 
   public static final String ISINSTANCE = "isinstance";
+  public static final String ASSERT_IS_INSTANCE = "assertIsInstance";
 
   public static final String DOCFORMAT = "__docformat__";
 

python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -10,6 +10,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator;
 import com.jetbrains.python.psi.impl.PyConstantExpressionEvaluator;
@@ -76,6 +77,11 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     else {
       super.visitPyCallExpression(node);
     }
+    if (node.isCalleeText(PyNames.ASSERT_IS_INSTANCE)) {
+      final PyTypeAssertionEvaluator assertionEvaluator = new PyTypeAssertionEvaluator();
+      node.accept(assertionEvaluator);
+      InstructionBuilder.addAssertInstructions(myBuilder, assertionEvaluator);
+    }
   }
 
   @Override

python/src/com/jetbrains/python/codeInsight/controlflow/PyTypeAssertionEvaluator.java

@@ -36,7 +36,7 @@ public class PyTypeAssertionEvaluator extends PyRecursiveElementVisitor {
 
   @Override
   public void visitPyCallExpression(PyCallExpression node) {
-    if (node.isCalleeText(PyNames.ISINSTANCE)) {
+    if (node.isCalleeText(PyNames.ISINSTANCE) || node.isCalleeText(PyNames.ASSERT_IS_INSTANCE)) {
       final PyExpression[] args = node.getArguments();
       if (args.length == 2 && args[0] instanceof PyReferenceExpression) {
         final PyReferenceExpression target = (PyReferenceExpression)args[0];

python/testSrc/com/jetbrains/python/PyTypeTest.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python;
 
+import com.jetbrains.python.documentation.PythonDocumentationProvider;
 import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyExpression;
@@ -262,11 +263,29 @@ public class PyTypeTest extends PyLightFixtureTestCase {
            "        expr = x");
   }
 
+  // PY-4383
+  public void testAssertIsInstance() {
+    doTest("int",
+           "from unittest import TestCase\n" +
+           "\n" +
+           "class Test1(TestCase):\n" +
+           "    def test_1(self, c):\n" +
+           "        x = 1 if c else 'foo'\n" +
+           "        self.assertIsInstance(x, int)\n" +
+           "        expr = x\n");
+  }
+
   private PyExpression parseExpr(String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     return myFixture.findElementByText("expr", PyExpression.class);
   }
 
+  private static String msg(PyType expected, PyType actual, TypeEvalContext context) {
+    return String.format("Expected: %s, actual: %s",
+                         PythonDocumentationProvider.getTypeName(expected, context),
+                         PythonDocumentationProvider.getTypeName(actual, context));
+  }
+
   private void doTest(final String expectedType, final String text) {
     PyExpression expr = parseExpr(text);
     TypeEvalContext context = TypeEvalContext.slow().withTracing();
@@ -274,7 +293,7 @@ public class PyTypeTest extends PyLightFixtureTestCase {
     PyType expected = PyTypeParser.getTypeByName(expr, expectedType);
     if (expected != null) {
       assertNotNull(context.printTrace(), actual);
-      assertTrue(PyTypeChecker.match(expected, actual, context));
+      assertTrue(msg(expected, actual, context), PyTypeChecker.match(expected, actual, context));
     }
   }
 }