Support for typing.NewType (PY-21302)

python/psi-api/src/com/jetbrains/python/psi/stubs/PyTypingNewTypeStub.java

@@ -0,0 +1,14 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.psi.stubs;
+
+import com.jetbrains.python.psi.impl.stubs.CustomTargetExpressionStub;
+import org.jetbrains.annotations.NotNull;
+
+public interface PyTypingNewTypeStub extends CustomTargetExpressionStub {
+
+  @NotNull
+  String getName();
+
+  @NotNull
+  String getClassType();
+}

python/src/META-INF/python-core-common.xml

@@ -708,6 +708,7 @@
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyNamedTupleStubType"/>
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType"/>
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyDataclassFieldStubType"/>
+    <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyTypingNewTypeStubType"/>
 
     <dialectsTokenSetContributor implementation="com.jetbrains.python.PythonTokenSetContributor"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>

python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -9,6 +9,7 @@ import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.ResolveResult;
 import com.intellij.psi.impl.source.resolve.FileContextUtil;
 import com.intellij.psi.util.CachedValueProvider;
 import com.intellij.psi.util.CachedValuesManager;
@@ -29,6 +30,7 @@ import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterType
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyCallExpressionHelper;
+import com.jetbrains.python.psi.impl.PyCallExpressionNavigator;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.stubs.PyClassElementType;
 import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType;
@@ -38,6 +40,8 @@ import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.search.PySuperMethodsSearch;
 import com.jetbrains.python.psi.stubs.PyClassStub;
+import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.stubs.PyTypingNewTypeStub;
 import com.jetbrains.python.psi.types.*;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
@@ -67,6 +71,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   public static final String PROTOCOL_EXT = "typing_extensions.Protocol";
   public static final String TYPE = "typing.Type";
   public static final String ANY = "typing.Any";
+  public static final String NEW_TYPE = "typing.NewType";
   public static final String CALLABLE = "typing.Callable";
   private static final String LIST = "typing.List";
   private static final String DICT = "typing.Dict";
@@ -82,7 +87,6 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   private static final String UNION = "typing.Union";
   private static final String OPTIONAL = "typing.Optional";
   private static final String CLASSVAR = "typing.ClassVar";
-  public static final String NEW_TYPE = "typing.NewType";
 
   public static final String NAMEDTUPLE_SIMPLE = "NamedTuple";
 
@@ -171,7 +175,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
         return createTypingCallableType(referenceExpression);
       }
     }
-    return null;
+
+    return getNewTypeForReference(referenceExpression, context);
   }
 
   @Override
@@ -458,6 +463,75 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
       return getOpenFunctionCallType(function, (PyCallExpression)callSite, LanguageLevel.forElement(callSite), context);
     }
 
+    if (callSite instanceof PyCallExpression && NEW_TYPE.equals(functionQName)) {
+      return Ref.create(getNewTypeForCallExpression((PyCallExpression)callSite, context));
+    }
+
+    return null;
+  }
+
+  @Nullable
+  private static PyType getNewTypeForReference(@NotNull PyReferenceExpression referenceExpression, @NotNull TypeEvalContext context) {
+    final PyCallExpression callee = PyCallExpressionNavigator.getPyCallExpressionByCallee(referenceExpression);
+    if (callee == null) {
+      return null;
+    }
+    final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
+    final ResolveResult[] resolveResults = referenceExpression.getReference(resolveContext).multiResolve(false);
+
+    for (PsiElement element : PyUtil.filterTopPriorityResults(resolveResults)) {
+      if (element instanceof PyTargetExpression) {
+        final PyType typeForTarget = getNewTypeCreationForTarget((PyTargetExpression)element, context);
+        if (typeForTarget != null) {
+          return typeForTarget;
+        }
+      }
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PyType getNewTypeCreationForTarget(@NotNull PyTargetExpression referenceTarget, @NotNull TypeEvalContext context) {
+    final PyTargetExpressionStub stub = referenceTarget.getStub();
+    if (stub != null) {
+      final PyTypingNewTypeStub customStub = stub.getCustomStub(PyTypingNewTypeStub.class);
+      if (customStub != null) {
+        final String newTypeName = customStub.getName();
+        final String classTypeStr = customStub.getClassType();
+        if (newTypeName != null && classTypeStr != null) {
+          final PyType type = Ref.deref(getStringBasedType(classTypeStr, referenceTarget, context));
+          if (type instanceof PyClassType) {
+            return new PyTypingNewType((PyClassType)type, true, newTypeName);
+          }
+        }
+      }
+    }
+    else {
+      final PyExpression value = referenceTarget.findAssignedValue();
+      if (value instanceof PyCallExpression) {
+        return getNewTypeForCallExpression(((PyCallExpression)value), context);
+      }
+    }
+    return null;
+  }
+
+  @Nullable
+  public static PyType getNewTypeForCallExpression(@NotNull PyCallExpression callExpression, @NotNull TypeEvalContext context) {
+    if (PyTypingNewType.Companion.isTypingNewType(callExpression)) {
+      final String className = PyResolveUtil.resolveFirstStrArgument(callExpression);
+      if (className != null) {
+        PyExpression secondArg = PyPsiUtils.flattenParens(callExpression.getArgument(1, PyExpression.class));
+        if (secondArg != null) {
+          final Ref<PyType> argType = getType(secondArg, new Context(context));
+          if (argType != null) {
+            final PyType type = argType.get();
+            if (type instanceof PyClassType) {
+              return new PyTypingNewType((PyClassType)type, true, className);
+            }
+          }
+        }
+      }
+    }
     return null;
   }
 
@@ -483,6 +557,11 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
         return ((PyInstantiableType)collection).toClass();
       }
 
+      final PyType newType = getNewTypeCreationForTarget(target, context);
+      if (newType != null) {
+        return newType;
+      }
+
       final Ref<PyType> annotatedType = getTypeFromTargetExpressionAnnotation(target, context);
       if (annotatedType != null) {
         return annotatedType.get();

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 68;
+    return 69;
   }
 
   @Nullable

python/src/com/jetbrains/python/psi/impl/stubs/PyNamedTupleStubImpl.java

@@ -78,7 +78,7 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub {
     final Pair<QualifiedName, NamedTupleModule> calleeNameAndModule = getCalleeNameAndNTModule(calleeReference);
 
     if (calleeNameAndModule != null) {
-      final String name = resolveTupleName(expression);
+      final String name = PyResolveUtil.resolveFirstStrArgument(expression);
 
       if (name == null) {
         return null;
@@ -160,33 +160,6 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub {
     return getImportedCalleeNameAndNTModule(referenceExpression);
   }
 
-  @Nullable
-  private static String resolveTupleName(@NotNull PyCallExpression callExpression) {
-    // SUPPORTED CASES:
-
-    // name = "Point"
-    // Point = namedtuple(name, ...)
-
-    // Point = namedtuple("Point", ...)
-
-    // Point = namedtuple(("Point"), ...)
-
-    // name = "Point"
-    // Point = NamedTuple(name, ...)
-
-    // Point = NamedTuple("Point", ...)
-
-    // Point = NamedTuple(("Point"), ...)
-
-    final PyExpression nameExpression = PyPsiUtils.flattenParens(callExpression.getArgument(0, PyExpression.class));
-
-    if (nameExpression instanceof PyReferenceExpression) {
-      return PyPsiUtils.strValue(fullResolveLocally((PyReferenceExpression)nameExpression));
-    }
-
-    return PyPsiUtils.strValue(nameExpression);
-  }
-
   @Nullable
   private static LinkedHashMap<String, Optional<String>> resolveTupleFields(@NotNull PyCallExpression callExpression, @NotNull NamedTupleModule module) {
     switch (module) {
@@ -315,23 +288,6 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub {
     return qualifiedName != null && name.equals(qualifiedName.toString());
   }
 
-  @Nullable
-  private static PyExpression fullResolveLocally(@NotNull PyReferenceExpression referenceExpression) {
-    for (PsiElement element : PyResolveUtil.resolveLocally(referenceExpression)) {
-      if (element instanceof PyTargetExpression) {
-        final PyExpression assignedValue = ((PyTargetExpression)element).findAssignedValue();
-
-        if (assignedValue instanceof PyReferenceExpression) {
-          return fullResolveLocally((PyReferenceExpression)assignedValue);
-        }
-
-        return assignedValue;
-      }
-    }
-
-    return null;
-  }
-
   @Nullable
   private static LinkedHashMap<String, Optional<String>> resolveCollectionsNTFields(@NotNull PyCallExpression callExpression) {
     // SUPPORTED CASES:
@@ -350,7 +306,7 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub {
     final PyExpression fields = PyPsiUtils.flattenParens(callExpression.getArgument(1, PyExpression.class));
 
     final PyExpression resolvedFields = fields instanceof PyReferenceExpression
-                                        ? fullResolveLocally((PyReferenceExpression)fields)
+                                        ? PyResolveUtil.fullResolveLocally((PyReferenceExpression)fields)
                                         : fields;
 
     final Collector<String, ?, LinkedHashMap<String, Optional<String>>> toFieldsOfUnknownType =
@@ -385,7 +341,7 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub {
       return getTypingNTFieldsFromKwArguments(Arrays.asList(arguments).subList(1, arguments.length));
     } else {
       final PyExpression resolvedFields = secondArgument instanceof PyReferenceExpression
-                                          ? fullResolveLocally((PyReferenceExpression)secondArgument)
+                                          ? PyResolveUtil.fullResolveLocally((PyReferenceExpression)secondArgument)
                                           : secondArgument;
       if (!(resolvedFields instanceof PySequenceExpression)) return null;
 

python/src/com/jetbrains/python/psi/impl/stubs/PyTypingNewTypeStubImpl.kt

@@ -0,0 +1,66 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.psi.PyCallExpression
+import com.jetbrains.python.psi.PyExpression
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.impl.PyPsiUtils
+import com.jetbrains.python.psi.resolve.PyResolveUtil
+import com.jetbrains.python.psi.stubs.PyTypingNewTypeStub
+import com.jetbrains.python.psi.types.PyTypingNewType
+import java.io.IOException
+
+class PyTypingNewTypeStubImpl private constructor(val qualifiedName: String,
+                                                  private val baseClassName: String) : PyTypingNewTypeStub {
+
+  override fun serialize(stream: StubOutputStream) {
+    stream.writeName(qualifiedName)
+    stream.writeName(baseClassName)
+  }
+
+  override fun getCalleeName(): QualifiedName? {
+    return QualifiedName.fromComponents(qualifiedName)
+  }
+
+  override fun getTypeClass(): Class<out CustomTargetExpressionStubType<out CustomTargetExpressionStub>> {
+    return PyTypingNewTypeStubType::class.java
+  }
+
+  companion object {
+    fun create(expression: PyTargetExpression?): PyTypingNewTypeStub? {
+      val callExpression = expression?.findAssignedValue() as? PyCallExpression ?: return null
+
+      if (PyTypingNewType.isTypingNewType(callExpression)) {
+        val newTypeName = PyResolveUtil.resolveFirstStrArgument(callExpression)
+        if (newTypeName != null) {
+          val secondArgument = PyPsiUtils.flattenParens(callExpression.getArgument(1, PyExpression::class.java))
+          if (secondArgument != null) {
+            return PyTypingNewTypeStubImpl(newTypeName, secondArgument.text)
+          }
+        }
+      }
+      return null
+    }
+
+    @Throws(IOException::class)
+    fun deserialize(stream: StubInputStream?): PyTypingNewTypeStub? {
+      if (stream != null) {
+        val typeName = stream.readName()
+        val classType = stream.readName()
+        if (typeName != null && classType != null) {
+          return PyTypingNewTypeStubImpl(typeName.string, classType.string)
+        }
+      }
+      return null
+    }
+  }
+
+  override fun getName() = qualifiedName
+
+  override fun getClassType() = baseClassName
+
+
+}

python/src/com/jetbrains/python/psi/impl/stubs/PyTypingNewTypeStubType.kt

@@ -0,0 +1,19 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.stubs.PyTypingNewTypeStub
+
+
+class PyTypingNewTypeStubType : CustomTargetExpressionStubType<PyTypingNewTypeStub>() {
+
+  override fun createStub(psi: PyTargetExpression?): PyTypingNewTypeStub? {
+    return PyTypingNewTypeStubImpl.create(psi)
+  }
+
+  override fun deserializeStub(stream: StubInputStream?): PyTypingNewTypeStub? {
+    return PyTypingNewTypeStubImpl.deserialize(stream)
+  }
+
+}

python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java

@@ -177,4 +177,48 @@ public class PyResolveUtil {
         }));
     return PyUtil.filterTopPriorityResults(result.toArray(RatedResolveResult[]::new));
   }
+
+  @Nullable
+  public static String resolveFirstStrArgument(@NotNull PyCallExpression callExpression) {
+    // SUPPORTED CASES:
+
+    // name = "Point"
+    // Point = namedtuple(name, ...)
+
+    // Point = namedtuple("Point", ...)
+
+    // Point = namedtuple(("Point"), ...)
+
+    // name = "Point"
+    // Point = NamedTuple(name, ...)
+
+    // Point = NamedTuple("Point", ...)
+
+    // Point = NamedTuple(("Point"), ...)
+
+    final PyExpression nameExpression = PyPsiUtils.flattenParens(callExpression.getArgument(0, PyExpression.class));
+
+    if (nameExpression instanceof PyReferenceExpression) {
+      return PyPsiUtils.strValue(fullResolveLocally((PyReferenceExpression)nameExpression));
+    }
+
+    return PyPsiUtils.strValue(nameExpression);
+  }
+
+  @Nullable
+  public static PyExpression fullResolveLocally(@NotNull PyReferenceExpression referenceExpression) {
+    for (PsiElement element : resolveLocally(referenceExpression)) {
+      if (element instanceof PyTargetExpression) {
+        final PyExpression assignedValue = ((PyTargetExpression)element).findAssignedValue();
+
+        if (assignedValue instanceof PyReferenceExpression) {
+          return fullResolveLocally((PyReferenceExpression)assignedValue);
+        }
+
+        return assignedValue;
+      }
+    }
+
+    return null;
+  }
 }

python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java

@@ -467,7 +467,7 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
 
   @NotNull
   @Override
-  public final List<PyClassLikeType> getAncestorTypes(@NotNull final TypeEvalContext context) {
+  public List<PyClassLikeType> getAncestorTypes(@NotNull final TypeEvalContext context) {
     return myClass.getAncestorTypes(context);
   }
 

python/src/com/jetbrains/python/psi/types/PyTypeChecker.java

@@ -259,6 +259,9 @@ public class PyTypeChecker {
                matchGenerics((PyCollectionType)expected, actual, context, substitutions, recursive, matching);
       }
       else if (matchClasses(superClass, subClass, context)) {
+        if (expectedClassType instanceof PyTypingNewType && !expectedClassType.equals(actualClassType) && superClass.equals(subClass)) {
+          return actualClassType.getAncestorTypes(context).contains(expectedClassType);
+        }
         return true;
       }
       if (expected.equals(actual)) {

python/src/com/jetbrains/python/psi/types/PyTypingNewType.kt

@@ -0,0 +1,71 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.psi.types
+
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.resolve.PyResolveUtil
+
+data class PyTypingNewType(internal val classType: PyClassType, internal val isDefinition: Boolean, internal val myName: String?) : PyClassTypeImpl(
+  classType.pyClass, isDefinition) {
+
+  override fun getName() = myName
+
+  override fun getCallType(context: TypeEvalContext, callSite: PyCallSiteExpression): PyType? {
+    return PyTypingNewType(classType, false, name)
+  }
+
+  override fun toClass(): PyClassLikeType {
+    return if (isDefinition) this else PyTypingNewType(classType, true, name)
+  }
+
+  override fun toInstance(): PyClassType {
+    return if (isDefinition) PyTypingNewType(classType, false, name) else this
+  }
+
+  override fun isBuiltin() = false
+
+  override fun isCallable() = classType.isCallable || isDefinition
+
+  override fun toString() = "TypingNewType: " + myName
+
+  override fun getParameters(context: TypeEvalContext): List<PyCallableParameter>? {
+    return if (isCallable) {
+      listOf(PyCallableParameterImpl.nonPsi(null, classType.toInstance(), null))
+    }
+    else {
+      null
+    }
+  }
+
+  override fun getSuperClassTypes(context: TypeEvalContext): List<PyClassLikeType> = listOf(classType)
+
+  override fun getAncestorTypes(context: TypeEvalContext): List<PyClassLikeType> {
+    return listOf(classType) + classType.getAncestorTypes(context)
+  }
+
+  companion object {
+    private fun getImportedQualifiedName(referenceExpression: PyReferenceExpression): QualifiedName? {
+      val qualifier = referenceExpression.qualifier
+      if (qualifier is PyReferenceExpression) {
+        PyResolveUtil.resolveLocally(qualifier)
+          .filterIsInstance<PyImportElement>()
+          .firstOrNull { return it.importedQName?.append(referenceExpression.name) }
+      }
+      for (element in PyResolveUtil.resolveLocally(referenceExpression)) {
+        if (element is PyImportElement) {
+          val importStatement = element.containingImportStatement
+          if (importStatement is PyFromImportStatement) {
+            return importStatement.importSourceQName?.append(element.importedQName)
+          }
+        }
+      }
+      return null
+    }
+
+    fun isTypingNewType(callExpression: PyCallExpression): Boolean {
+      val calleeReference = callExpression.callee as? PyReferenceExpression ?: return false
+      return getImportedQualifiedName(calleeReference) == QualifiedName.fromDottedString(PyTypingTypeProvider.NEW_TYPE)
+    }
+  }
+}

python/testData/inspections/PyTypeCheckerInspection/InitializingNewType.py

@@ -0,0 +1,12 @@
+from typing import NewType, Dict
+
+UserId = NewType("UserId", int)
+
+a = UserId(42)
+b = UserId(<warning descr="Expected type 'int', got 'str' instead">"John"</warning>)
+
+KeyValue = NewType("KeyValue", Dict[str, int])
+
+KeyValue({"key": 13})
+KeyValue(<warning descr="Expected type 'Dict[str, int]', got 'int' instead">42</warning>)
+KeyValue(<warning descr="Expected type 'Dict[str, int]', got 'Dict[str, str]' instead">{"key1": "key2"}</warning>)

python/testData/inspections/PyTypeCheckerInspection/NewTypeAsParameter.py

@@ -0,0 +1,11 @@
+from typing import NewType
+
+UserId = NewType("UserId", int)
+
+def get_user(user: UserId) -> str:
+    pass
+
+
+get_user(UserId(5))
+get_user(<warning descr="Expected type 'UserId', got 'str' instead">"John"</warning>)
+get_user(<warning descr="Expected type 'UserId', got 'int' instead">4</warning>)

python/testData/inspections/PyTypeCheckerInspection/NewTypeInheritance.py

@@ -0,0 +1,31 @@
+from typing import NewType
+
+UserId = NewType("UserId", int)
+NewId = NewType("NewId", UserId)
+ChildNewId = NewType("ChildNewId", NewId)
+
+def get_user_super(user: UserId) -> str:
+    pass
+
+def get_user_child(user: NewId) -> str:
+    pass
+
+def get_user_child_new(user: ChildNewId):
+    pass
+
+
+user = UserId(12)
+new_id = NewId(user)
+child_new_id = ChildNewId(new_id)
+
+get_user_super(user)
+get_user_super(new_id)
+get_user_super(child_new_id)
+
+get_user_child(<warning descr="Expected type 'NewId', got 'UserId' instead">user</warning>)
+get_user_child(new_id)
+get_user_child(child_new_id)
+
+get_user_child_new(<warning descr="Expected type 'ChildNewId', got 'UserId' instead">user</warning>)
+get_user_child_new(<warning descr="Expected type 'ChildNewId', got 'NewId' instead">new_id</warning>)
+get_user_child_new(child_new_id)

python/testData/stubs/typingNewType/new_type.py

@@ -0,0 +1,3 @@
+from typing import NewType
+
+UserId = NewType("UserId", int)

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -24,10 +24,7 @@ import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyFileImpl;
 import com.jetbrains.python.psi.resolve.PyResolveImportUtil;
 import com.jetbrains.python.psi.stubs.*;
-import com.jetbrains.python.psi.types.PyCallableType;
-import com.jetbrains.python.psi.types.PyClassType;
-import com.jetbrains.python.psi.types.PyType;
-import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.psi.types.*;
 import com.jetbrains.python.toolbox.Maybe;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -972,4 +969,27 @@ public class PyStubsTest extends PyTestCase {
       }
     );
   }
+
+  // PY-27398
+  public void testTypingNewType() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> {
+        final PyFile file = getTestFile("typingNewType/new_type.py");
+
+        final PyTargetExpression type = file.findTopLevelAttribute("UserId");
+        final PyTypingNewTypeStub stub = type.getStub().getCustomStub(PyTypingNewTypeStub.class);
+
+        assertNotNull(stub);
+        assertTrue("UserId".equals(stub.getName()));
+        assertTrue("int".equals(stub.getClassType()));
+
+        final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
+        final PyType typeDef = context.getType(type);
+
+        assertTrue(typeDef instanceof PyTypingNewType);
+        assertNotParsed(file);
+      }
+    );
+  }
 }

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

@@ -8,10 +8,7 @@ import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
-import com.jetbrains.python.psi.types.PyClassLikeType;
-import com.jetbrains.python.psi.types.PyClassType;
-import com.jetbrains.python.psi.types.PyType;
-import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.psi.types.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -3078,6 +3075,92 @@ public class PyTypeTest extends PyTestCase {
     );
   }
 
+  // PY-21302
+  public void testNewTypeReferenceTarget() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> {
+        final PyExpression definition = parseExpr("from typing import NewType\n" +
+                                                  "UserId = NewType('UserId', int)\n" +
+                                                  "expr = UserId");
+
+        for (TypeEvalContext context : getTypeEvalContexts(definition)) {
+          assertInstanceOf(context.getType(definition), PyTypingNewType.class);
+        }
+
+        final PyExpression instance = parseExpr("from typing import NewType\n" +
+                                                "UserId = NewType('UserId', int)\n" +
+                                                "expr = UserId(12)");
+
+        for (TypeEvalContext context : getTypeEvalContexts(instance)) {
+          assertInstanceOf(context.getType(instance), PyTypingNewType.class);
+        }
+      }
+    );
+  }
+
+  // PY-21302
+  public void testNewType() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("UserId",
+                   "from typing import NewType\n" +
+                   "UserId = NewType('UserId', int)\n" +
+                   "expr = UserId(12)")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("Type[UserId]",
+                   "from typing import Dict, NewType\n" +
+                   "UserId = NewType('UserId', Dict[int, str])\n" +
+                   "expr = UserId\n")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("(a: UserId) -> str",
+                   "from typing import Dict, NewType\n" +
+                   "UserId = NewType('UserId', int)\n" +
+                   "def foo(a: UserId) -> str\n" +
+                   "    pass\n" +
+                   "expr = foo\n")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("UserId",
+                   "from typing import NewType as nt\n" +
+                   "UserId = nt('UserId', int)\n" +
+                   "expr = UserId(12)\n")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("UserId",
+                   "import typing\n" +
+                   "UserId = typing.NewType('UserId', int)\n" +
+                   "expr = UserId(12)\n")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("UserId",
+                   "import typing as t\n" +
+                   "UserId = t.NewType('UserId', int)\n" +
+                   "expr = UserId(12)\n")
+    );
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("SuperId",
+                   "from typing import NewType\n" +
+                   "UserId = NewType('UserId', int)\n" +
+                   "SuperId = NewType('SuperId', UserId)\n" +
+                   "expr = SuperId(UserId(12))\n")
+    );
+  }
+
   // PY-26992
   public void testImportedOrderedDict() {
     doTest("OrderedDict[str, str]",

python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java

@@ -246,6 +246,21 @@ public class Py3TypeCheckerInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-21302
+  public void testInitializingNewType() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
+  // PY-21302
+  public void testNewTypeAsParameter() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
+  // PY-21302
+  public void testNewTypeInheritance() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
   // PY-24287
   public void testPromotingBytearrayToBytes() {
     doTest();