Fixed type inference for type() calls (PY-7058)

python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java

@@ -110,6 +110,25 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
             return superCallType.value();
           }
         }
+        if ("type".equals(callee.getText())) {
+          final PyExpression[] args = getArguments();
+          if (args.length == 1) {
+            final PyExpression arg = args[0];
+            final PyType argType = arg.getType(context);
+            if (argType instanceof PyClassType) {
+              final PyClassType classType = (PyClassType)argType;
+              if (!classType.isDefinition()) {
+                final PyClass cls = classType.getPyClass();
+                if (cls != null) {
+                  return cls.getType(context);
+                }
+              }
+            }
+            else {
+              return null;
+            }
+          }
+        }
         // normal cases
         final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
         ResolveResult[] targets = ((PyReferenceExpression)callee).getReference(resolveContext).multiResolve(false);

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

@@ -385,6 +385,40 @@ public class PyTypeTest extends PyTestCase {
     assertNull(actual);
   }
 
+  // PY-7058
+  public void testReturnTypeOfTypeForInstance() {
+    PyExpression expr = parseExpr("class C(object):\n" +
+                                  "    pass\n" +
+                                  "\n" +
+                                  "x = C()\n" +
+                                  "expr = type(x)\n");
+    TypeEvalContext context = TypeEvalContext.slow().withTracing();
+    PyType type = expr.getType(context);
+    assertInstanceOf(type, PyClassType.class);
+    assertTrue("Got instance type instead of class type", ((PyClassType)type).isDefinition());
+  }
+
+  // PY-7058
+  public void testReturnTypeOfTypeForClass() {
+    PyExpression expr = parseExpr("class C(object):\n" +
+                                  "    pass\n" +
+                                  "\n" +
+                                  "expr = type(C)\n");
+    TypeEvalContext context = TypeEvalContext.slow().withTracing();
+    PyType type = expr.getType(context);
+    assertInstanceOf(type, PyClassType.class);
+    assertEquals(type.getName(), "type");
+  }
+
+  // PY-7058
+  public void testReturnTypeOfTypeForUnknown() {
+    PyExpression expr = parseExpr("def f(x):\n" +
+                                  "    expr = type(x)\n");
+    TypeEvalContext context = TypeEvalContext.slow().withTracing();
+    PyType type = expr.getType(context);
+    assertNull(type);
+  }
+
   private PyExpression parseExpr(String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     return myFixture.findElementByText("expr", PyExpression.class);