[Findit] Adding support for splitting try job data by trybot

appengine/findit/util_scripts/remote_queries/try_job_data_metrics.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Pulls historical try job metadata from Findit and prints a report."""
 
 import argparse
 from collections import defaultdict
 import datetime
+import json
 import numpy
 import os
 import sys
 
 _REMOTE_API_DIR = os.path.join(os.path.dirname(__file__), os.path.pardir)
 sys.path.insert(1, _REMOTE_API_DIR)
 
 import remote_api
 
 from model.wf_try_job_data import WfTryJobData
@@ skipping 27 matching lines @@
         'ios': 'ios',
     }
 
     for os_name, platform in os_map.iteritems():
       if os_name in builder_name:
         return platform
 
     return 'unknown'
 
 
+def _GetTrybotName(buildbucket_response):
+  # First try parameters_json.
+  parameters = json.loads(buildbucket_response.get('parameters_json', '{}'))
+  trybot = parameters.get('builder_name')
+
+  if not trybot:
+    # Fallback to response_details_json.
+    response_details = json.loads(
+        buildbucket_response.get('response_details_json', '{}'))
+    properties = response_details.get('properties', {})
+    trybot = properties.get('buildername')
+
+    if not trybot:
+      trybot = 'unknown'
+
+  return trybot
+
+
 def _GetAverageOfNumbersInList(numbers):
   """Returns a float average of numbers or NOT_AVAILABLE if numbers is empty."""
   return (float(sum(numbers)) / len(numbers)) if numbers else NOT_AVAILABLE
 
 
 def _FormatDigits(number):
   """Formats number into a 2-digit float, or NOT_AVAILABLE."""
   if isinstance(number, float):
     return float('%.2f' % number)
   return NOT_AVAILABLE
@@ skipping 331 matching lines @@
 
     if not master_name or not builder_name:
       continue
 
     platform = _GetOSPlatformName(master_name, builder_name)
     categorized_data_dict[platform].append(try_job_data)
 
   return categorized_data_dict
 
 
+def _SplitListByTrybot(try_job_data_list):
+  categorized_data_dict = defaultdict(list)
+
+  for try_job_data in try_job_data_list:
+    if not try_job_data.last_buildbucket_response:
+      continue
+
+    trybot = _GetTrybotName(try_job_data.last_buildbucket_response)
+    categorized_data_dict[trybot].append(try_job_data)
+
+  return categorized_data_dict
+
+
 def SplitListByOption(try_job_data_list, option):
   """Takes a WfTryJobData list and separates it into a dict based on arg.
 
   Args:
     try_job_data_list: A list of WfTryJobData entities.
     option: An option with which to split the data by.
 
   Returns:
     A dict where the keys are how the data is separated based on arg and the
     values are the corresponding lists of data.
   """
 
-  if option == 't':  # Try job type.
-    return _SplitListByTryJobType(try_job_data_list)
-  elif option == 'm':  # Main waterfall master.
-    return _SplitListByMaster(try_job_data_list)
-  elif option == 'b':  # Main waterfall builder.
+  if option == 'b':  # Main waterfall builder.
     return _SplitListByBuilder(try_job_data_list)
-  elif option == 'r':  # Whether or not heuristic results are included.
-    return _SplitListByHeuristicResults(try_job_data_list)
   elif option == 'c':  # Whether or not compile targets are included.
     return _SplitListByCompileTargets(try_job_data_list)
   elif option == 'e':  # Whether or not try jobs with errors should be counted.
     return _SplitListByError(try_job_data_list)
+  elif option == 'm':  # Main waterfall master.
+    return _SplitListByMaster(try_job_data_list)
   elif option == 'p':  # Split by OS platform.
     return _SplitListByPlatform(try_job_data_list)
+  elif option == 'r':  # Whether or not heuristic results are included.
+    return _SplitListByHeuristicResults(try_job_data_list)
+  elif option == 't':  # Try job type.
+    return _SplitListByTryJobType(try_job_data_list)
+  elif 'trybot' in option:  # Split by trybot.
+    return _SplitListByTrybot(try_job_data_list)
 
   # Unsupported flag, bail out without modification.
   return try_job_data_list
 
 
 def SplitStructByOption(try_job_data_struct, option):
   if isinstance(try_job_data_struct, list):
     try_job_data_struct = SplitListByOption(try_job_data_struct, option)
   elif isinstance(try_job_data_struct, dict):
     for key, struct in try_job_data_struct.iteritems():
@@ skipping 18 matching lines @@
                             'errors detected'))
   parser.add_argument('-m', action='store_true',
                       help='group try job data by master')
   parser.add_argument('-p', action='store_true',
                       help='group try job data by platform')
   parser.add_argument('-r', action='store_true',
                       help=('group try job data by those with and without '
                             'heuristic results'))
   parser.add_argument('-t', action='store_true',
                       help='group try job data by type (compile, test)')
+  parser.add_argument('--trybot', action='store_true',
+                      help='group try job data by trybot')
 
   args_dict = vars(parser.parse_args())
 
   # Preserve order from original command.
   ordered_args = []
 
   for original_arg in command_line_args:
-    parsed_arg = original_arg[1:]
+    parsed_arg = original_arg.lstrip('-')
     if args_dict[parsed_arg]:
       ordered_args.append(parsed_arg)
 
   return ordered_args
 
 
 if __name__ == '__main__':
   # Set up the Remote API to use services on the live App Engine.
   remote_api.EnableRemoteApi(app_id='findit-for-me')
 
   START_DATE = datetime.datetime(2016, 5, 1)
   END_DATE = datetime.datetime(2016, 6, 23)
 
   try_job_data_query = WfTryJobData.query(
       WfTryJobData.request_time >= START_DATE,
       WfTryJobData.request_time < END_DATE)
   categorized_data = try_job_data_query.fetch()
 
   args = GetArgsInOrder()
   for arg in args:
     categorized_data = SplitStructByOption(categorized_data, arg)
 
   # TODO(lijeffrey): Display data in an html page instead of printing.
   PrettyPrint(categorized_data, START_DATE, END_DATE)