Bug 1302786 - Support filtering output from task-graph generation; r?dustin draft bug-1302786
authorHammad Akhtar <hammad13060@iiitd.ac.in>
Tue, 15 Nov 2016 07:35:59 +0530
branchbug-1302786
changeset 438797 4f980b5b6d3b4fed6969fc50fe65b0582827d61b
parent 437797 5d1c54f240d077b6b8e141fd44b52a3b694a5dfd
child 724615 ad0c2019e3193accba1f856b00e8ca1efe5025b6
push id35816
push userhammad13060@iiitd.ac.in
push dateTue, 15 Nov 2016 02:07:16 +0000
reviewersdustin
bugs1302786
milestone52.0a1
Bug 1302786 - Support filtering output from task-graph generation; r?dustin MozReview-Commit-ID: 181zcy6rxw9
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,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,45 +210,68 @@ 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'))
-            show_method(tg)
+
+            filtered_tg_dict = self.filter_tasks(tg, options['tasks_regex'])
+
+            tg = TaskGraph.from_json(filtered_tg_dict)
+            
+            show_method(tg[1])
         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 filter_tasks(self, taskgraph, tasksregex):
+
+        if not tasksregex:
+            return taskgraph.to_json()
+
+
+        tasks = taskgraph.to_json()
+        filteredtasks = {}
+        regexprogram = re.compile(tasksregex)
+        for key in tasks:
+            if regexprogram.match(tasks[key]["label"]):
+                filteredtasks[key] = tasks[key]
+        return filteredtasks
+
+
+
+
 
 @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")