Bug 1725404 - Refactor '--target-kind' / '--output-file' logic to use standalone taskgraph's implementation, r=bhearsum
authorAndrew Halberstadt <ahal@mozilla.com>
Tue, 17 Aug 2021 13:35:12 +0000
changeset 589109 e2712f329b8c59d1d919ac4304e7c87960bc7953
parent 589108 217b1fc04f06a62f532ae79276de98a9206be542
child 589110 61127b8f451c707a3522fc5eb4dca266eb5343d7
push id148153
push userahalberstadt@mozilla.com
push dateTue, 17 Aug 2021 13:37:38 +0000
treeherderautoland@61127b8f451c [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersbhearsum
bugs1725404
milestone93.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1725404 - Refactor '--target-kind' / '--output-file' logic to use standalone taskgraph's implementation, r=bhearsum This syncs in the latest 'main.py' from standalone taskgraph and refactors mach_commands.py to use it. Differential Revision: https://phabricator.services.mozilla.com/D122778
taskcluster/mach_commands.py
taskcluster/taskgraph/main.py
--- a/taskcluster/mach_commands.py
+++ b/taskcluster/mach_commands.py
@@ -24,17 +24,17 @@ from mach.decorators import (
     CommandProvider,
     SettingsProvider,
     SubCommand,
 )
 from mozbuild.base import MachCommandBase
 
 from taskgraph.main import (
     commands as taskgraph_commands,
-    get_filtered_taskgraph,
+    format_taskgraph,
 )
 
 logger = logging.getLogger("taskcluster")
 
 
 @SettingsProvider
 class TaskgraphConfig(object):
     @classmethod
@@ -68,31 +68,16 @@ class ShowTaskGraphSubCommand(SubCommand
     """A SubCommand with TaskGraph-specific arguments"""
 
     def __call__(self, func):
         name = func.__name__.replace("_", "-").split("-", 1)[1]
         args = taskgraph_commands[name].func.args
 
         extra_args = [
             (
-                ["--target-kind"],
-                {
-                    "default": None,
-                    "help": "only return tasks that are of the given kind, "
-                    "or their dependencies.",
-                },
-            ),
-            (
-                ["-o", "--output-file"],
-                {
-                    "default": None,
-                    "help": "file path to store generated output.",
-                },
-            ),
-            (
                 ["--diff"],
                 {
                     "const": "default",
                     "nargs": "?",
                     "default": None,
                     "help": "Generate and diff the current taskgraph against another revision. "
                     "Without args the base revision will be used. A revision specifier such as "
                     "the hash or `.~1` (hg) or `HEAD~1` (git) can be used as well.",
@@ -206,45 +191,51 @@ class MachCommands(MachCommandBase):
         by dependencies: for example, a binary must be built before it is tested,
         and that build may further depend on various toolchains, libraries, etc.
         """
 
     @ShowTaskGraphSubCommand(
         "taskgraph", "tasks", description="Show all tasks in the taskgraph"
     )
     def taskgraph_tasks(self, command_context, **options):
-        return self.show_taskgraph(command_context, "full_task_set", options)
+        options["graph_attr"] = "full_task_set"
+        return self.show_taskgraph(command_context, options)
 
     @ShowTaskGraphSubCommand("taskgraph", "full", description="Show the full taskgraph")
     def taskgraph_full(self, command_context, **options):
-        return self.show_taskgraph(command_context, "full_task_graph", options)
+        options["graph_attr"] = "full_task_graph"
+        return self.show_taskgraph(command_context, options)
 
     @ShowTaskGraphSubCommand(
         "taskgraph", "target", description="Show the target task set"
     )
     def taskgraph_target(self, command_context, **options):
-        return self.show_taskgraph(command_context, "target_task_set", options)
+        options["graph_attr"] = "target_task_set"
+        return self.show_taskgraph(command_context, options)
 
     @ShowTaskGraphSubCommand(
         "taskgraph", "target-graph", description="Show the target taskgraph"
     )
     def taskgraph_target_graph(self, command_context, **options):
-        return self.show_taskgraph(command_context, "target_task_graph", options)
+        options["graph_attr"] = "target_task_graph"
+        return self.show_taskgraph(command_context, options)
 
     @ShowTaskGraphSubCommand(
         "taskgraph", "optimized", description="Show the optimized taskgraph"
     )
     def taskgraph_optimized(self, command_context, **options):
-        return self.show_taskgraph(command_context, "optimized_task_graph", options)
+        options["graph_attr"] = "optimized_task_graph"
+        return self.show_taskgraph(command_context, options)
 
     @ShowTaskGraphSubCommand(
         "taskgraph", "morphed", description="Show the morphed taskgraph"
     )
     def taskgraph_morphed(self, command_context, **options):
-        return self.show_taskgraph(command_context, "morphed_task_graph", options)
+        options["graph_attr"] = "morphed_task_graph"
+        return self.show_taskgraph(command_context, options)
 
     @SubCommand("taskgraph", "actions", description="Write actions.json to stdout")
     @CommandArgument(
         "--root", "-r", help="root of the taskgraph definition relative to topsrcdir"
     )
     @CommandArgument(
         "--quiet", "-q", action="store_true", help="suppress all logging output"
     )
@@ -356,42 +347,45 @@ class MachCommands(MachCommandBase):
                 level=level,
                 write_interval=old.formatter.write_interval,
                 write_times=old.formatter.write_times,
             )
 
         # all of the taskgraph logging is unstructured logging
         command_context.log_manager.enable_unstructured()
 
-    def show_taskgraph(self, command_context, graph_attr, options):
+    def show_taskgraph(self, command_context, options):
         self.setup_logging(
             command_context, quiet=options["quiet"], verbose=options["verbose"]
         )
         vcs = None
         base_out = ""
         base_ref = None
         cur_ref = None
 
+        if not options["parameters"]:
+            options["parameters"] = "project=mozilla-central"
+
         if options["diff"]:
             from mozversioncontrol import get_repository_object
 
             vcs = get_repository_object(command_context.topsrcdir)
             with vcs:
                 if not vcs.working_directory_clean():
                     print("abort: can't diff taskgraph with dirty working directory")
                     return 1
 
                 # We want to return the working directory to the current state
                 # as best we can after we're done. In all known cases, using
                 # branch or bookmark (which are both available on the VCS object)
                 # as `branch` is preferable to a specific revision.
                 cur_ref = vcs.branch or vcs.head_ref[:12]
-            logger.info("Generating {} @ {}".format(graph_attr, cur_ref))
+            logger.info("Generating {} @ {}".format(options["graph_attr"], cur_ref))
 
-        out = self.format_taskgraph(graph_attr, options)
+        out = format_taskgraph(options)
 
         if options["diff"]:
             with vcs:
                 # Some transforms use global state for checks, so will fail
                 # when running taskgraph a second time in the same session.
                 # Reload all taskgraph modules to avoid this.
                 for mod in sys.modules.copy():
                     if mod.startswith("taskgraph"):
@@ -400,23 +394,27 @@ class MachCommands(MachCommandBase):
                 if options["diff"] == "default":
                     base_ref = vcs.base_ref
                 else:
                     base_ref = options["diff"]
 
                 try:
                     vcs.update(base_ref)
                     base_ref = vcs.head_ref[:12]
-                    logger.info("Generating {} @ {}".format(graph_attr, base_ref))
-                    base_out = self.format_taskgraph(graph_attr, options)
+                    logger.info(
+                        "Generating {} @ {}".format(options["graph_attr"], base_ref)
+                    )
+                    base_out = format_taskgraph(options)
                 finally:
                     vcs.update(cur_ref)
 
             diffcmd = command_context._mach_context.settings["taskgraph"]["diffcmd"]
-            diffcmd = diffcmd.format(attr=graph_attr, base=base_ref, cur=cur_ref)
+            diffcmd = diffcmd.format(
+                attr=options["graph_attr"], base=base_ref, cur=cur_ref
+            )
 
             with tempfile.NamedTemporaryFile(mode="w") as base:
                 base.write(base_out)
 
                 with tempfile.NamedTemporaryFile(mode="w") as cur:
                     cur.write(out)
                     try:
                         out = subprocess.run(
@@ -442,60 +440,16 @@ class MachCommands(MachCommandBase):
             fh = open(fh, "w")
         print(out, file=fh)
         if options["format"] != "json":
             logger.info(
                 "If you were expecting differences in task bodies "
                 'you should pass "-J"\n'
             )
 
-    def format_taskgraph(self, graph_attr, options):
-        import taskgraph
-        import taskgraph.generator
-        import taskgraph.parameters
-
-        if not options["parameters"]:
-            options["parameters"] = "project=mozilla-central"
-
-        if options["fast"]:
-            taskgraph.fast = True
-
-        try:
-            parameters = taskgraph.parameters.parameters_loader(
-                options["parameters"],
-                overrides={"target-kind": options.get("target_kind")},
-                strict=False,
-            )
-
-            tgg = taskgraph.generator.TaskGraphGenerator(
-                root_dir=options.get("root"),
-                parameters=parameters,
-            )
-
-            tg = getattr(tgg, graph_attr)
-            tg = get_filtered_taskgraph(tg, options["tasks_regex"])
-
-            format_method = getattr(
-                self, "format_taskgraph_" + (options["format"] or "labels")
-            )
-            return format_method(tg)
-        except Exception:
-            traceback.print_exc()
-            sys.exit(1)
-
-    def format_taskgraph_labels(self, taskgraph):
-        return "\n".join(
-            taskgraph.tasks[index].label for index in taskgraph.graph.visit_postorder()
-        )
-
-    def format_taskgraph_json(self, taskgraph):
-        return json.dumps(
-            taskgraph.to_json(), sort_keys=True, indent=2, separators=(",", ": ")
-        )
-
     def show_actions(self, command_context, options):
         import taskgraph
         import taskgraph.actions
         import taskgraph.generator
         import taskgraph.parameters
 
         try:
             self.setup_logging(
--- a/taskcluster/taskgraph/main.py
+++ b/taskcluster/taskgraph/main.py
@@ -35,31 +35,30 @@ def argument(*args, **kwargs):
         if not hasattr(func, "args"):
             func.args = []
         func.args.append((args, kwargs))
         return func
 
     return decorator
 
 
-def show_taskgraph_labels(taskgraph):
-    for index in taskgraph.graph.visit_postorder():
-        print(taskgraph.tasks[index].label)
-
-
-def show_taskgraph_json(taskgraph):
-    print(
-        json.dumps(
-            taskgraph.to_json(), sort_keys=True, indent=2, separators=(",", ": ")
-        )
+def format_taskgraph_labels(taskgraph):
+    return "\n".join(
+        taskgraph.tasks[index].label for index in taskgraph.graph.visit_postorder()
     )
 
 
-def show_taskgraph_yaml(taskgraph):
-    print(yaml.safe_dump(taskgraph.to_json(), default_flow_style=False))
+def format_taskgraph_json(taskgraph):
+    return json.dumps(
+        taskgraph.to_json(), sort_keys=True, indent=2, separators=(",", ": ")
+    )
+
+
+def format_taskgraph_yaml(taskgraph):
+    return yaml.safe_dump(taskgraph.to_json(), default_flow_style=False)
 
 
 def get_filtered_taskgraph(taskgraph, tasksregex):
     """
     Filter all the tasks on basis of a regular expression
     and returns a new TaskGraph object
     """
     from taskgraph.graph import Graph
@@ -81,23 +80,52 @@ def get_filtered_taskgraph(taskgraph, ta
                 if regexprogram.match(dep):
                     filterededges.add((key, dep, depname))
     filtered_taskgraph = TaskGraph(
         filteredtasks, Graph(set(filteredtasks), filterededges)
     )
     return filtered_taskgraph
 
 
-SHOW_METHODS = {
-    "labels": show_taskgraph_labels,
-    "json": show_taskgraph_json,
-    "yaml": show_taskgraph_yaml,
+FORMAT_METHODS = {
+    "labels": format_taskgraph_labels,
+    "json": format_taskgraph_json,
+    "yaml": format_taskgraph_yaml,
 }
 
 
+def format_taskgraph(options):
+    import taskgraph
+    import taskgraph.parameters
+    import taskgraph.generator
+
+    if options["fast"]:
+        taskgraph.fast = True
+
+    try:
+        parameters = taskgraph.parameters.parameters_loader(
+            options["parameters"],
+            overrides={"target-kind": options.get("target_kind")},
+            strict=False,
+        )
+
+        tgg = taskgraph.generator.TaskGraphGenerator(
+            root_dir=options.get("root"), parameters=parameters
+        )
+
+        tg = getattr(tgg, options["graph_attr"])
+        tg = get_filtered_taskgraph(tg, options["tasks_regex"])
+
+        format_method = FORMAT_METHODS[options["format"] or "labels"]
+        return format_method(tg)
+    except Exception:
+        traceback.print_exc()
+        sys.exit(1)
+
+
 @command(
     "tasks",
     help="Show all tasks in the taskgraph.",
     defaults={"graph_attr": "full_task_set"},
 )
 @command(
     "full", help="Show the full taskgraph.", defaults={"graph_attr": "full_task_graph"}
 )
@@ -157,58 +185,50 @@ SHOW_METHODS = {
     help="parameters file (.yml or .json; see " "`taskcluster/docs/parameters.rst`)`",
 )
 @argument(
     "--no-optimize",
     dest="optimize",
     action="store_false",
     default="true",
     help="do not remove tasks from the graph that are found in the "
-    "index (a.k.a. ptimize the graph)",
+    "index (a.k.a. optimize the graph)",
+)
+@argument(
+    "-o",
+    "--output-file",
+    default=None,
+    help="file path to store generated output.",
 )
 @argument(
     "--tasks-regex",
     "--tasks",
     default=None,
     help="only return tasks with labels matching this regular " "expression.",
 )
 @argument(
+    "--target-kind",
+    default=None,
+    help="only return tasks that are of the given kind, or their dependencies.",
+)
+@argument(
     "-F",
     "--fast",
     dest="fast",
     default=False,
     action="store_true",
     help="enable fast task generation for local debugging.",
 )
 def show_taskgraph(options):
-    import taskgraph.parameters
-    import taskgraph.target_tasks
-    import taskgraph.generator
-    import taskgraph
-
-    if options["fast"]:
-        taskgraph.fast = True
-
-    try:
-        parameters = taskgraph.parameters.parameters_loader(
-            options["parameters"], strict=False
-        )
+    out = format_taskgraph(options)
 
-        tgg = taskgraph.generator.TaskGraphGenerator(
-            root_dir=options.get("root"), parameters=parameters
-        )
-
-        tg = getattr(tgg, options["graph_attr"])
-
-        show_method = SHOW_METHODS.get(options["format"] or "labels")
-        tg = get_filtered_taskgraph(tg, options["tasks_regex"])
-        show_method(tg)
-    except Exception:
-        traceback.print_exc()
-        sys.exit(1)
+    fh = options["output_file"]
+    if fh:
+        fh = open(fh, "w")
+    print(out, file=fh)
 
 
 @command("build-image", help="Build a Docker image")
 @argument("image_name", help="Name of the image to build")
 @argument(
     "-t", "--tag", help="tag that the image should be built as.", metavar="name:tag"
 )
 @argument(