From d5673dceef43dd7c08cc56007fc9d8adbab26124 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 19 Jun 2025 21:09:44 +0300
Subject: [PATCH] [python] Make overlooked PyElementVisitor methods for
 statements delegate to visitPyStatement

(cherry picked from commit f6853d0dfc5ddcc7b40bfccead0ad779328fc751)

IJ-CR-172556

GitOrigin-RevId: b7d424930b1f4752ea69041a4ab53229b0392f5d
---
 .../src/com/jetbrains/python/psi/PyElementVisitor.java      | 6 +++---
 .../codeInsight/controlflow/PyControlFlowBuilder.java       | 4 +++-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java b/python/python-psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
index cc7444aa65e0..8e08bd26b225 100644
--- a/python/python-psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
+++ b/python/python-psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
@@ -203,11 +203,11 @@ public class PyElementVisitor extends PsiElementVisitor {
   }
 
   public void visitPyFunction(@NotNull PyFunction node) {
-    visitPyElement(node);
+    visitPyStatement(node);
   }
 
   public void visitPyClass(@NotNull PyClass node) {
-    visitPyElement(node);
+    visitPyStatement(node);
   }
 
   public void visitPyFile(@NotNull PyFile node) {
@@ -271,7 +271,7 @@ public class PyElementVisitor extends PsiElementVisitor {
   }
 
   public void visitPyAssertStatement(@NotNull PyAssertStatement node) {
-    visitPyElement(node);
+    visitPyStatement(node);
   }
 
   public void visitPyPassStatement(@NotNull PyPassStatement node) {
diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
index b1afc6f90ec6..6440b2a88217 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
@@ -949,8 +949,10 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
   @Override
   public void visitPyAssertStatement(final @NotNull PyAssertStatement node) {
     myBuilder.startNode(node);
-    super.visitPyAssertStatement(node);
     final PyExpression[] args = node.getArguments();
+    for (PyExpression arg : args) {
+      arg.accept(this);
+    }
     // assert False
     if (args.length >= 1) {
       if (!PyEvaluator.evaluateAsBooleanNoResolve(args[0], true)) {