Raise a warning on Final inside a loop (PEP 591) (PY-34945)

GitOrigin-RevId: 0685b853f2971af91539f4776024b2ca682c62d8

python/src/com/jetbrains/python/inspections/PyFinalInspection.kt

@@ -3,8 +3,10 @@ package com.jetbrains.python.inspections
 
 import com.intellij.codeInspection.LocalInspectionToolSession
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.jetbrains.python.PyNames
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
 import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterTypeList
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.*
@@ -126,6 +128,16 @@ class PyFinalInspection : PyInspection() {
       checkFinalIsOuterMost(node)
     }
 
+    override fun visitPyForStatement(node: PyForStatement) {
+      super.visitPyForStatement(node)
+      checkFinalInsideLoop(node)
+    }
+
+    override fun visitPyWhileStatement(node: PyWhileStatement) {
+      super.visitPyWhileStatement(node)
+      checkFinalInsideLoop(node)
+    }
+
     private fun getClassLevelFinalsAndInitAttributes(cls: PyClass): Pair<Map<String?, PyTargetExpression>, Map<String, PyTargetExpression>> {
       val classLevelFinals = mutableMapOf<String?, PyTargetExpression>()
       cls.classAttributes.forEach { if (isFinal(it)) classLevelFinals[it.name] = it }
@@ -313,6 +325,24 @@ class PyFinalInspection : PyInspection() {
       }
     }
 
+    private fun checkFinalInsideLoop(loop: PyLoopStatement) {
+      loop.acceptChildren(
+        object : PyRecursiveElementVisitor() {
+          override fun visitElement(element: PsiElement) {
+            if (element !is ScopeOwner) super.visitElement(element)
+          }
+
+          override fun visitPyForStatement(node: PyForStatement) {}
+
+          override fun visitPyWhileStatement(node: PyWhileStatement) {}
+
+          override fun visitPyTargetExpression(node: PyTargetExpression) {
+            if (isFinal(node)) registerProblem(node, "'Final' could not be used inside a loop")
+          }
+        }
+      )
+    }
+
     private fun isFinal(decoratable: PyDecoratable) = isFinal(decoratable, myTypeEvalContext)
 
     private fun <T> isFinal(node: T): Boolean where T : PyAnnotationOwner, T : PyTypeCommentOwner {

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

@@ -500,6 +500,27 @@ public class PyFinalInspectionTest extends PyInspectionTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doMultiFileTest);
   }
 
+  // PY-34945
+  public void testFinalInsideLoop() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTestByText("from typing_extensions import Final\n" +
+                         "\n" +
+                         "for i in undefined:\n" +
+                         "    if undefined:\n" +
+                         "        <warning descr=\"'Final' could not be used inside a loop\">x</warning>: Final[int] = 1\n" +
+                         "while undefined:\n" +
+                         "    <warning descr=\"'Final' could not be used inside a loop\">y</warning>: Final[str] = '1'\n" +
+                         "    \n" +
+                         "def foo():\n" +
+                         "    for i in undefined:\n" +
+                         "        if undefined:\n" +
+                         "            <warning descr=\"'Final' could not be used inside a loop\">x</warning>: Final[int] = 1\n" +
+                         "    while undefined:\n" +
+                         "        <warning descr=\"'Final' could not be used inside a loop\">y</warning>: Final[str] = '1'")
+    );
+  }
+
   @NotNull
   @Override
   protected Class<? extends PyInspection> getInspectionClass() {