From 2a962ffd63ae0ae8c45f53c3cdd88663be03b794 Mon Sep 17 00:00:00 2001
From: Charles Brunet <charles.brunet@optelgroup.com>
Date: Fri, 22 Sep 2023 15:42:34 -0400
Subject: [PATCH] visit full ast for AstConditionLevel

Use the new `FullAstVisitor` for `AstConditionLevel`. This will allow
proper formatting of `if / else / endif` and `foreach / endforeach`
blocks.
---
 mesonbuild/ast/postprocess.py | 41 +++++++++++++++++++++--------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/mesonbuild/ast/postprocess.py b/mesonbuild/ast/postprocess.py
index e9f251345..7acfc03a0 100644
--- a/mesonbuild/ast/postprocess.py
+++ b/mesonbuild/ast/postprocess.py
@@ -5,7 +5,7 @@
 # or an interpreter-based tool
 from __future__ import annotations
 
-from .visitor import AstVisitor
+from .visitor import AstVisitor, FullAstVisitor
 import typing as T
 
 if T.TYPE_CHECKING:
@@ -78,32 +78,41 @@ class AstIDGenerator(AstVisitor):
         node.ast_id = name + '#' + str(self.counter[name])
         self.counter[name] += 1
 
-class AstConditionLevel(AstVisitor):
+class AstConditionLevel(FullAstVisitor):
     def __init__(self) -> None:
         self.condition_level = 0
 
-    def visit_default_func(self, node: mparser.BaseNode) -> None:
+    def enter_node(self, node: mparser.BaseNode) -> None:
         node.condition_level = self.condition_level
 
     def visit_ForeachClauseNode(self, node: mparser.ForeachClauseNode) -> None:
-        self.visit_default_func(node)
-        self.condition_level += 1
+        self.enter_node(node)
+        node.foreach_.accept(self)
+        for varname in node.varnames:
+            varname.accept(self)
+        for comma in node.commas:
+            comma.accept(self)
+        node.colon.accept(self)
         node.items.accept(self)
+        self.condition_level += 1
         node.block.accept(self)
         self.condition_level -= 1
-
-    def visit_IfClauseNode(self, node: mparser.IfClauseNode) -> None:
-        self.visit_default_func(node)
-        for i in node.ifs:
-            i.accept(self)
-        if node.elseblock:
-            self.condition_level += 1
-            node.elseblock.accept(self)
-            self.condition_level -= 1
+        node.endforeach.accept(self)
+        self.exit_node(node)
 
     def visit_IfNode(self, node: mparser.IfNode) -> None:
-        self.visit_default_func(node)
-        self.condition_level += 1
+        self.enter_node(node)
+        node.if_.accept(self)
         node.condition.accept(self)
+        self.condition_level += 1
+        node.block.accept(self)
+        self.condition_level -= 1
+        self.exit_node(node)
+
+    def visit_ElseNode(self, node: mparser.ElseNode) -> None:
+        self.enter_node(node)
+        node.else_.accept(self)
+        self.condition_level += 1
         node.block.accept(self)
         self.condition_level -= 1
+        self.exit_node(node)