diff --git a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib.py b/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib.py index e3277177f401..8c7804c8b1a2 100644 --- a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib.py +++ b/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} diff --git a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib_test.py b/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib_test.py index 49099e7f0b5b..3d840548c874 100644 --- a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/lib_test.py +++ b/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) diff --git a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe.py b/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe.py index f4d31397756c..60c7a61f698e 100644 --- a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe.py +++ b/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 ), diff --git a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe_test.py b/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe_test.py index ab03fa2db66d..f2b29deb19eb 100644 --- a/pkgs/by-name/fl/flatten-references-graph/src/flatten_references_graph/pipe_test.py +++ b/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)