Bug 1302786 - Support filtering output from task-graph generation; r=dustin
authorHammad Akhtar <hammad13060@iiitd.ac.in>
Fri, 18 Nov 2016 04:38:28 +0530
changeset 370507 d607a23fc3dc64956e68b588e53772b9aef48f4f
parent 370506 144ca2b74fd2c4f37fde99e293b2f4104e87b9e1
child 370508 add1700c24ca4e854f9c4e66bf433cc8c6696f2d
push id1419
push userjlund@mozilla.com
push dateMon, 10 Apr 2017 20:44:07 +0000
treeherdermozilla-release@5e6801b73ef6 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersdustin
bugs1302786
milestone53.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 1302786 - Support filtering output from task-graph generation; r=dustin MozReview-Commit-ID: 7efKKRgvX3r
taskcluster/mach_commands.py
--- a/taskcluster/mach_commands.py
+++ b/taskcluster/mach_commands.py
@@ -6,16 +6,17 @@
 
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 import json
 import logging
 import sys
 import traceback
+import re
 
 from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
     SubCommand,
 )
 
@@ -44,16 +45,18 @@ class ShowTaskGraphSubCommand(SubCommand
                             help="Output the label for each task in the task graph (default)"),
             CommandArgument('--parameters', '-p', required=True,
                             help="parameters file (.yml or .json; see "
                                  "`taskcluster/docs/parameters.rst`)`"),
             CommandArgument('--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. optimize the graph)"),
+            CommandArgument('--tasks-regex', '--tasks', default=None)
+
         ]
         for arg in args:
             after = arg(after)
         return after
 
 
 @CommandProvider
 class MachCommands(MachCommandBase):
@@ -224,29 +227,55 @@ class MachCommands(MachCommandBase):
             tgg = taskgraph.generator.TaskGraphGenerator(
                 root_dir=options['root'],
                 parameters=parameters,
                 target_tasks_method=target_tasks_method)
 
             tg = getattr(tgg, graph_attr)
 
             show_method = getattr(self, 'show_taskgraph_' + (options['format'] or 'labels'))
+            tg = self.get_filtered_taskgraph(tg, options["tasks_regex"])
             show_method(tg)
         except Exception:
             traceback.print_exc()
             sys.exit(1)
 
     def show_taskgraph_labels(self, taskgraph):
         for label in taskgraph.graph.visit_postorder():
             print(label)
 
     def show_taskgraph_json(self, taskgraph):
         print(json.dumps(taskgraph.to_json(),
               sort_keys=True, indent=2, separators=(',', ': ')))
 
+    def get_filtered_taskgraph(self, taskgraph, tasksregex):
+        from taskgraph.graph import Graph
+        from taskgraph.taskgraph import TaskGraph
+        """
+        This class method filters all the tasks on basis of a regular expression
+        and returns a new TaskGraph object
+        """
+        # return original taskgraph if no regular expression is passed
+        if not tasksregex:
+            return taskgraph
+        named_links_dict = taskgraph.graph.named_links_dict()
+        filteredtasks = {}
+        filterededges = set()
+        regexprogram = re.compile(tasksregex)
+
+        for key in taskgraph.graph.visit_postorder():
+            task = taskgraph.tasks[key]
+            if regexprogram.match(task.label):
+                filteredtasks[key] = task
+                for depname, dep in named_links_dict[key].iteritems():
+                    if regexprogram.match(dep):
+                        filterededges.add((key, dep, depname))
+        filtered_taskgraph = TaskGraph(filteredtasks, Graph(set(filteredtasks), filterededges))
+        return filtered_taskgraph
+
 
 @CommandProvider
 class TaskClusterImagesProvider(object):
     @Command('taskcluster-load-image', category="ci",
              description="Load a pre-built Docker image")
     @CommandArgument('--task-id',
                      help="Load the image at public/image.tar.zst in this task,"
                           "rather than searching the index")