flatten-references-graph: add reorder DSL primitive for explicit flatten order (#544325)

pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib.py

@@ -327,3 +327,25 @@ def remove_paths(paths, graph):
 @curry
 def reverse(iterator):
     return reversed(list(iterator))
+
+
+@curry
+def reorder(keys, graph):
+    """Reorder the keys of a dict of graphs, controlling the traversal order
+    when followed by flatten(). Keys absent from the graph are ignored;
+    keys present in the graph but not in `keys` are appended at the end in
+    their original insertion order.
+
+    Example: after subcomponent_out the result is {"main": ..., "rest": ...}
+    and flatten() yields main before rest. Use reorder(["rest", "main"]) to
+    put rest-layers first without the double-reverse workaround.
+
+    Also works for split_paths which returns {"main", "common", "rest"}:
+    reorder(["rest", "common", "main"]) puts shared deps (common) before the
+    split-off paths (main).
+    """
+    if not isinstance(graph, dict):
+        return graph
+    ordered = {k: graph[k] for k in keys if k in graph}
+    remainder = {k: v for k, v in graph.items() if k not in ordered}
+    return {**ordered, **remainder}

pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib_test.py

@@ -10,7 +10,8 @@ from .lib import (
     limit_layers,
     pick_keys,
     references_graph_to_igraph,
-    reference_graph_node_keys_to_keep
+    reference_graph_node_keys_to_keep,
+    reorder
 )
 
 if __name__ == "__main__":
@@ -197,3 +198,45 @@ class TestLib(unittest.TestCase, th.CustomAssertions):
         )
 
         self.assertGraphEqual(graph, result_list[0])
+
+
+class TestReorder(unittest.TestCase):
+
+    def test_reorder_two_keys_swaps_main_and_rest(self):
+        g1 = directed_graph([], ["A"])
+        g2 = directed_graph([], ["B"])
+        d = {"main": g1, "rest": g2}
+        result = reorder(["rest", "main"], d)
+        self.assertEqual(list(result.keys()), ["rest", "main"])
+        self.assertIs(result["rest"], g2)
+        self.assertIs(result["main"], g1)
+
+    def test_reorder_three_keys_split_paths_order(self):
+        g_main   = directed_graph([], ["A"])
+        g_common = directed_graph([], ["B"])
+        g_rest   = directed_graph([], ["C"])
+        d = {"main": g_main, "common": g_common, "rest": g_rest}
+        result = reorder(["rest", "common", "main"], d)
+        self.assertEqual(list(result.keys()), ["rest", "common", "main"])
+
+    def test_reorder_missing_keys_are_ignored(self):
+        g1 = directed_graph([], ["A"])
+        g2 = directed_graph([], ["B"])
+        d = {"main": g1, "rest": g2}
+        # "common" is not in dict — should be silently skipped
+        result = reorder(["rest", "common", "main"], d)
+        self.assertEqual(list(result.keys()), ["rest", "main"])
+
+    def test_reorder_extra_keys_appended_in_original_order(self):
+        g1 = directed_graph([], ["A"])
+        g2 = directed_graph([], ["B"])
+        g3 = directed_graph([], ["C"])
+        g4 = directed_graph([], ["D"])
+        d = {"main": g1, "rest": g2, "extra1": g3, "extra2": g4}
+        result = reorder(["rest", "main"], d)
+        self.assertEqual(list(result.keys()), ["rest", "main", "extra1", "extra2"])
+
+    def test_reorder_non_dict_passthrough(self):
+        graphs = [directed_graph([], ["A"]), directed_graph([], ["B"])]
+        result = reorder(["rest", "main"], graphs)
+        self.assertIs(result, graphs)

pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe.py

@@ -20,7 +20,8 @@ funcs = tlz.merge(
             "split_every",
             "limit_layers",
             "remove_paths",
-            "reverse"
+            "reverse",
+            "reorder"
         ],
         lib
     ),

pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe_test.py

@@ -151,3 +151,62 @@ class Test(
                 ([], ["A", "B"])
             ]
         )
+
+    def test_reorder_rest_before_main(self):
+        # Verify that ["reorder" ["rest" "main"]] puts the rest-graph layers
+        # before the main (split-off) layers, equivalent to the double-reverse
+        # workaround but expressed directly.
+        #
+        # Graph: Root -> A -> B -> D
+        #                  -> C
+        # Split off B and its deps; rest should come first.
+        graph = directed_graph(
+            [
+                ("Root", "A"),
+                ("A", "B"),
+                ("A", "C"),
+                ("B", "D"),
+            ]
+        )
+
+        # Without reorder: main (B+D) comes before rest (Root, A, C)
+        result_without = list(pipe(
+            [
+                ["subcomponent_out", ["B"]],
+                ["flatten"],
+            ],
+            graph
+        ))
+        # main graph contains B and D
+        main_names = set(result_without[0].vs["name"])
+        self.assertIn("B", main_names)
+
+        # With reorder: rest (Root, A, C) comes before main (B+D)
+        result_with = list(pipe(
+            [
+                ["subcomponent_out", ["B"]],
+                ["reorder", ["rest", "main"]],
+                ["flatten"],
+            ],
+            graph
+        ))
+        rest_names = set(result_with[0].vs["name"])
+        self.assertIn("A", rest_names)
+        self.assertNotIn("B", rest_names)
+
+    def test_reorder_three_keys_split_paths(self):
+        # split_paths returns {"main", "common", "rest"} — verify reorder
+        # can express "rest, common, main" ordering for stable shared layers first.
+        graph = make_test_graph()
+        result = list(pipe(
+            [
+                ["split_paths", ["B"]],
+                ["reorder", ["rest", "common", "main"]],
+                ["flatten"],
+                ["limit_layers", 3],
+            ],
+            graph
+        ))
+        # rest should be first: contains Root1, A, C (no B or its exclusive deps)
+        first_names = set(result[0].vs["name"])
+        self.assertNotIn("B", first_names)