Bug 1302786 - taskgraph.taskgraph.TaskGraph statement removed bug-1302786-3
authorHammad Akhtar <hammad13060@iiitd.ac.in>
Thu, 17 Nov 2016 00:21:20 +0530
branchbug-1302786-3
changeset 894890 75b55c5a80f7f054a02fcea222f8a6c9840e9b36
parent 891032 5d1c54f240d077b6b8e141fd44b52a3b694a5dfd
child 894891 cb963f50aaebd3d80525ab0c277437680305decc
child 894914 c02a869744745c7fd7933a03b9bb2a5e3e1be715
child 896234 3c63c4d3ada7c6d6c3e1eed1b56d141232098003
push id155183
push userhammad13060@iiitd.ac.in
push dateWed, 16 Nov 2016 18:56:39 +0000
treeherdertry@cb963f50aaeb [default view] [failures only]
bugs1302786
milestone52.0a1
Bug 1302786 - taskgraph.taskgraph.TaskGraph statement removed MozReview-Commit-ID: KcHXheRc8gR
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.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 in this task,"
                           "rather than searching the index")