From 064419d1e9227016a7a26c49d91d00ce8f03be26 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

GitOrigin-RevId: f6853d0dfc5ddcc7b40bfccead0ad779328fc751
---
 .../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 f260c1079817..997f863adfc4 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)) {