Bug 1302786 - Support filtering output from task-graph generation; r?dustin bug-1302786-2
authorHammad Akhtar <hammad13060@iiitd.ac.in>
Tue, 15 Nov 2016 23:23:00 +0530
branchbug-1302786-2
changeset 894615 9d857c040a96dd05b9065e255cd4a9dbbfbab2bf
parent 891032 5d1c54f240d077b6b8e141fd44b52a3b694a5dfd
child 894616 dc79b735fafb78317f4b7ff409995a43a193b0c2
push id155136
push userhammad13060@iiitd.ac.in
push dateWed, 16 Nov 2016 15:55:27 +0000
treeherdertry@471eaa285bcd [default view] [failures only]
reviewersdustin
bugs1302786
milestone52.0a1
Bug 1302786 - Support filtering output from task-graph generation; r?dustin MozReview-Commit-ID: Bw9KFym8AFn
taskcluster/mach_commands.py
taskcluster/taskgraph/taskgraph.py
--- a/taskcluster/mach_commands.py
+++ b/taskcluster/mach_commands.py
@@ -44,16 +44,17 @@ 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):
@@ -208,32 +209,33 @@ class MachCommands(MachCommandBase):
 
         # all of the taskgraph logging is unstructured logging
         self.log_manager.enable_unstructured()
 
     def show_taskgraph(self, graph_attr, options):
         import taskgraph.parameters
         import taskgraph.target_tasks
         import taskgraph.generator
+        from taskgraph.taskgraph import TaskGraph
 
         try:
             self.setup_logging(quiet=options['quiet'], verbose=options['verbose'])
             parameters = taskgraph.parameters.load_parameters_file(options)
             parameters.check()
 
             target_tasks_method = parameters.get('target_tasks_method', 'all_tasks')
             target_tasks_method = taskgraph.target_tasks.get_method(target_tasks_method)
             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 = TaskGraph.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)
--- a/taskcluster/taskgraph/taskgraph.py
+++ b/taskcluster/taskgraph/taskgraph.py
@@ -1,15 +1,16 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 import os
+import re
 
 from .graph import Graph
 from .util.python_path import find_object
 
 TASKCLUSTER_QUEUE_URL = "https://queue.taskcluster.net/v1/task/"
 GECKO = os.path.realpath(os.path.join(__file__, '..', '..', '..'))
 
 
@@ -75,8 +76,32 @@ class TaskGraph(object):
             task_kind = find_object(implementation)
             tasks[key] = task_kind.from_json(value)
             if 'task_id' in value:
                 tasks[key].task_id = value['task_id']
             for depname, dep in value['dependencies'].iteritems():
                 edges.add((key, dep, depname))
         task_graph = cls(tasks, Graph(set(tasks), edges))
         return tasks, task_graph
+
+    @classmethod
+    def get_filtered_taskgraph(cls, taskgraph, tasksregex):
+        """
+        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.iteritems():
+                    if regexprogram.match(dep):
+                        filterededges.add((key, dep, depname))
+        filtered_taskgraph = cls(filteredtasks, Graph(set(filteredtasks), filterededges))
+        return filtered_taskgraph