Add prelude and epilogue (#6784)

### Problem

It would be nice to execute code before and after children are visited in graph traversals. The context of this change is the implementation of Dependency Banning (#6541), and the propagation of constraints through the dependency graph.

### Solution

Add optional prelude and epilogue function parameters to the dep. graph walk functions, that get executed right before visiting children. This should not incur any API or performance changes for existing code.

### Result

Calls like these are now possible:
```
a = self.make_target('a')
b = self.make_target('b', dependencies=[a])
names = []
self.build_graph.walk_transitive_dependency_graph(
    [b.address],
    work=lambda x: names.append(x.name),
    postorder=postorder,
    prelude=lambda x: names.append(x.name + "-pre"),
     epilogue=lambda x: names.append(x.name + "-epi"),
)
assertEqual(names, ['b', 'b-pre', 'a', 'a-pre', 'a-epi', 'b-epi'])
```

Blocks: #6541

src/python/pants/build_graph/build_graph.py

@@ -347,7 +347,9 @@ def walk_transitive_dependency_graph(self,
                                        work,
                                        predicate=None,
                                        postorder=False,
-                                       dep_predicate=None):
+                                       dep_predicate=None,
+                                       prelude=None,
+                                       epilogue=None):
     """Given a work function, walks the transitive dependency closure of `addresses` using DFS.
 
     :API: public
@@ -365,6 +367,16 @@ def walk_transitive_dependency_graph(self,
       the current target. If this parameter is not given, no dependencies will be filtered
       when traversing the closure. If it is given, when the predicate fails, the edge to the dependency
       will not be expanded.
+    :param function prelude: Function to run before any dependency expansion.
+      It takes the currently explored target as an argument.
+      If ``postorder`` is ``False``, it will run after the current node is visited.
+      It is not affected by ``dep_predicate``.
+      It is not run if ``predicate`` does not succeed.
+    :param function epilogue: Function to run after children nodes are visited.
+      It takes the currently explored target as an argument.
+      If ``postorder`` is ``True``, it runs before visiting the current node.
+      It is not affected by ``dep_predicate``.
+      It is not run if ``predicate`` is not passed.
     """
     walk = self._walk_factory(dep_predicate)
 
@@ -381,19 +393,31 @@ def _walk_rec(addr, level=0):
       if not postorder and walk.do_work_once(addr):
         work(target)
 
+      if prelude:
+        prelude(target)
+
       for dep_address in self._target_dependencies_by_address[addr]:
         if walk.expanded_or_worked(dep_address):
           continue
         if walk.dep_predicate(target, self._target_by_address[dep_address], level):
           _walk_rec(dep_address, level + 1)
 
+      if epilogue:
+        epilogue(target)
+
       if postorder and walk.do_work_once(addr):
         work(target)
 
     for address in addresses:
       _walk_rec(address)
 
-  def walk_transitive_dependee_graph(self, addresses, work, predicate=None, postorder=False):
+  def walk_transitive_dependee_graph(self,
+                                     addresses,
+                                     work,
+                                     predicate=None,
+                                     postorder=False,
+                                     prelude=None,
+                                     epilogue=None):
     """Identical to `walk_transitive_dependency_graph`, but walks dependees preorder (or postorder
     if the postorder parameter is True).
 
@@ -411,8 +435,12 @@ def _walk_rec(addr):
         if not predicate or predicate(target):
           if not postorder:
             work(target)
+          if prelude:
+            prelude(target)
           for dep_address in self._target_dependees_by_address[addr]:
             _walk_rec(dep_address)
+          if epilogue:
+            epilogue(target)
           if postorder:
             work(target)
     for address in addresses:

tests/python/pants_test/build_graph/test_build_graph.py

@@ -145,6 +145,46 @@ def assertDependeeWalk(target, results, postorder=False):
     assertDependencyWalk(a, [a, b, c, d, e])
     assertDependencyWalk(a, [c, d, b, e, a], postorder=True)
 
+  def test_dependency_walk_prelude_epilogue(self):
+    def assertDependencyWalkPreludeEpilogue(target, results, postorder=False):
+      names = []
+      self.build_graph.walk_transitive_dependency_graph([target.address],
+        lambda x: names.append(x.name),
+        postorder=postorder,
+        prelude=lambda x: names.append(x.name + "-pre"),
+        epilogue=lambda x: names.append(x.name + "-epi"),
+      )
+      self.assertEqual(results, names)
+
+    a = self.make_target('a')
+    b = self.make_target('b', dependencies=[a])
+
+    assertDependencyWalkPreludeEpilogue(a, ['a', 'a-pre', 'a-epi'])
+    assertDependencyWalkPreludeEpilogue(b, ['b', 'b-pre', 'a', 'a-pre', 'a-epi', 'b-epi'])
+
+    assertDependencyWalkPreludeEpilogue(a, ['a-pre', 'a-epi', 'a'], postorder=True)
+    assertDependencyWalkPreludeEpilogue(b, ['b-pre', 'a-pre', 'a-epi', 'a', 'b-epi', 'b'], postorder=True)
+
+  def test_dependee_walk_prelude_epilogue(self):
+    def assertDependeeWalkPreludeEpilogue(target, results, postorder=False):
+      names = []
+      self.build_graph.walk_transitive_dependee_graph([target.address],
+        lambda x: names.append(x.name),
+        postorder=postorder,
+        prelude=lambda x: names.append(x.name + "-pre"),
+        epilogue=lambda x: names.append(x.name + "-epi"),
+      )
+      self.assertEqual(results, names)
+
+    a = self.make_target('a')
+    b = self.make_target('b', dependencies=[a])
+
+    assertDependeeWalkPreludeEpilogue(a, ['a', 'a-pre', 'b', 'b-pre', 'b-epi', 'a-epi'])
+    assertDependeeWalkPreludeEpilogue(b, ['b', 'b-pre', 'b-epi'])
+
+    assertDependeeWalkPreludeEpilogue(a, ['a-pre', 'b-pre', 'b-epi', 'b', 'a-epi', 'a'], postorder=True)
+    assertDependeeWalkPreludeEpilogue(b, ['b-pre', 'b-epi', 'b'], postorder=True)
+
   def test_target_closure(self):
     a = self.make_target('a')
     self.assertEqual([a], a.closure())