From 158dc1aeab8760ec6e6afe2b43a72dbe4dbeccfd Mon Sep 17 00:00:00 2001
From: Petr <petr.golubev@jetbrains.com>
Date: Fri, 6 Dec 2024 13:56:33 +0100
Subject: [PATCH] PY-76842 Conformance test failure: aliases_newtype.py

NewType's type argument check.

GitOrigin-RevId: f4c2c6fc61b604458cb6182f1e4db7f8f53f2087
---
 .../PyNewTypeInspection.html                  | 19 ++++++++++-----
 .../resources/messages/PyPsiBundle.properties |  5 +++-
 .../python/inspections/PyNewTypeInspection.kt | 24 +++++++++++++++++++
 .../inspections/PyNewTypeInspectionTest.kt    | 23 ++++++++++++++++++
 4 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/python/python-psi-impl/resources/inspectionDescriptions/PyNewTypeInspection.html b/python/python-psi-impl/resources/inspectionDescriptions/PyNewTypeInspection.html
index b3e8f05dff56..8f917efdb93c 100644
--- a/python/python-psi-impl/resources/inspectionDescriptions/PyNewTypeInspection.html
+++ b/python/python-psi-impl/resources/inspectionDescriptions/PyNewTypeInspection.html
@@ -4,16 +4,23 @@
   Reports invalid usages of <a href="https://docs.python.org/3/library/typing.html#typing.NewType">NewType</a>.
 </p>
 <p>
-  <b>Example:</b>
+  <b>Examples:</b>
 </p>
 <pre><code>
-from typing import NewType
+  from typing import NewType
 
-Base = NewType("Base", str)
-A = NewType("B", int)  # Variable name 'A' does not match NewType name 'B'
+  InvalidName = NewType("Name", int)  # Variable name 'InvalidName' does not match NewType name 'Name'
+</code></pre>
+<pre><code>
+  from typing import Literal
 
-class Derived(Base):  # 'Base' cannot be subclassed
-    pass
+  InvalidType = NewType("InvalidType", Literal[1])  # NewType cannot be used with 'Literal[1]'
+</code></pre>
+<pre><code>
+  Base = NewType("Base", str)
+
+  class Derived(Base):  # 'Base' cannot be subclassed
+      pass
 </code></pre>
 </body>
 </html>
\ No newline at end of file
diff --git a/python/python-psi-impl/resources/messages/PyPsiBundle.properties b/python/python-psi-impl/resources/messages/PyPsiBundle.properties
index 02933c49ed54..ec0a2de176c4 100644
--- a/python/python-psi-impl/resources/messages/PyPsiBundle.properties
+++ b/python/python-psi-impl/resources/messages/PyPsiBundle.properties
@@ -1249,4 +1249,7 @@ INSP.enum.type.annotations.are.not.allowed.for.enum.members=Type annotations are
 #PyNewTypeInspection
 INSP.NAME.new.type=Invalid usage of NewType
 INSP.NAME.new.type.cannot.be.subclassed=''{0}'' cannot be subclassed
-INSP.NAME.new.type.variable.name.does.not.match.new.type.name=Variable name ''{0}'' does not match NewType name ''{1}''
\ No newline at end of file
+INSP.NAME.new.type.variable.name.does.not.match.new.type.name=Variable name ''{0}'' does not match NewType name ''{1}''
+INSP.NAME.new.type.expected.class=Expected class
+INSP.NAME.new.type.new.type.cannot.be.used.with=NewType cannot be used with ''{0}''
+INSP.NAME.new.type.new.type.cannot.be.used.with.protocol.classes=NewType cannot be used with protocol classes
\ No newline at end of file
diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyNewTypeInspection.kt b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyNewTypeInspection.kt
index ba017b7c6eb9..2d26edba2aec 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyNewTypeInspection.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyNewTypeInspection.kt
@@ -3,13 +3,20 @@ package com.jetbrains.python.inspections
 import com.intellij.codeInspection.LocalInspectionToolSession
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElementVisitor
 import com.jetbrains.python.PyPsiBundle
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
+import com.jetbrains.python.codeInsight.typing.isProtocol
 import com.jetbrains.python.psi.PyCallExpression
 import com.jetbrains.python.psi.PyClass
+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.types.PyClassType
+import com.jetbrains.python.psi.types.PyLiteralType
+import com.jetbrains.python.psi.types.PyTypedDictType
 import com.jetbrains.python.psi.types.PyTypingNewType
 
 class PyNewTypeInspection : PyInspection() {
@@ -30,6 +37,23 @@ class PyNewTypeInspection : PyInspection() {
             registerProblem(node.nameIdentifier,
                             PyPsiBundle.message("INSP.NAME.new.type.variable.name.does.not.match.new.type.name", targetName, newTypeName))
           }
+
+          val typeExpr = PyPsiUtils.flattenParens(assignedValue.getArgument(1, "tp", PyExpression::class.java))
+          if (typeExpr != null) {
+            val type = Ref.deref(PyTypingTypeProvider.getType(typeExpr, myTypeEvalContext))
+            if (type !is PyClassType) {
+              registerProblem(typeExpr, PyPsiBundle.message("INSP.NAME.new.type.expected.class"))
+            }
+            else if (type is PyLiteralType) {
+              registerProblem(typeExpr, PyPsiBundle.message("INSP.NAME.new.type.new.type.cannot.be.used.with", type.name))
+            }
+            else if (type is PyTypedDictType) {
+              registerProblem(typeExpr, PyPsiBundle.message("INSP.NAME.new.type.new.type.cannot.be.used.with", "TypedDict"))
+            }
+            else if (isProtocol(type, myTypeEvalContext)) {
+              registerProblem(typeExpr, PyPsiBundle.message("INSP.NAME.new.type.new.type.cannot.be.used.with.protocol.classes"))
+            }
+          }
         }
       }
 
diff --git a/python/testSrc/com/jetbrains/python/inspections/PyNewTypeInspectionTest.kt b/python/testSrc/com/jetbrains/python/inspections/PyNewTypeInspectionTest.kt
index 744ec63ba824..24e090e15bf3 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyNewTypeInspectionTest.kt
+++ b/python/testSrc/com/jetbrains/python/inspections/PyNewTypeInspectionTest.kt
@@ -33,5 +33,28 @@ class PyNewTypeInspectionTest : PyInspectionTestCase() {
     )
   }
 
+  fun testType() {
+    doTestByText(
+      """
+        from typing import NewType, Literal, TypedDict, Protocol
+        
+        class TD(TypedDict):
+            name: str
+        
+        class MyProtocol(Protocol):
+            pass
+        
+        NewType1 = NewType(tp=int, name="NewType1")
+        NewType2 = NewType("NewType2", NewType1)
+        NewType3 = NewType("NewType3", <warning descr="Expected class">int | str</warning>)
+        NewType4 = NewType("NewType4", <warning descr="Expected class">1</warning>)
+        NewType5 = NewType("NewType5", <warning descr="NewType cannot be used with 'Literal[1]'">Literal[1]</warning>)
+        NewType6 = NewType("NewType6", <warning descr="NewType cannot be used with 'TypedDict'">TD</warning>)
+        NewType7 = NewType("NewType7", <warning descr="NewType cannot be used with protocol classes">MyProtocol</warning>)
+        NewType8 = NewType("NewType8", list[int])
+      """.trimIndent()
+    )
+  }
+
   override fun getInspectionClass(): Class<out PyInspection> = PyNewTypeInspection::class.java
 }
\ No newline at end of file