src/tools/telemetry/telemetry/core/discover.py

   1 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import fnmatch
   6 import inspect
   7 import os
   8 import re
   9
  10 from telemetry import decorators
  11 from telemetry.core import camel_case
  12 from telemetry.core import util
  13 from telemetry.page import page_set
  14
  15
  16 @decorators.Cache
  17 def DiscoverModules(start_dir, top_level_dir, pattern='*'):
  18   """Discover all modules in |start_dir| which match |pattern|.
  19
  20   Args:
  21     start_dir: The directory to recursively search.
  22     top_level_dir: The top level of the package, for importing.
  23     pattern: Unix shell-style pattern for filtering the filenames to import.
  24
  25   Returns:
  26     list of modules.
  27   """
  28   modules = []
  29   for dir_path, _, filenames in os.walk(start_dir):
  30     for filename in filenames:
  31       # Filter out unwanted filenames.
  32       if filename.startswith('.') or filename.startswith('_'):
  33         continue
  34       if os.path.splitext(filename)[1] != '.py':
  35         continue
  36       if not fnmatch.fnmatch(filename, pattern):
  37         continue
  38
  39       # Find the module.
  40       module_rel_path = os.path.relpath(os.path.join(dir_path, filename),
  41                                         top_level_dir)
  42       module_name = re.sub(r'[/\\]', '.', os.path.splitext(module_rel_path)[0])
  43
  44       # Import the module.
  45       module = __import__(module_name, fromlist=[True])
  46
  47       modules.append(module)
  48   return modules
  49
  50
  51 # TODO(dtu): Normalize all discoverable classes to have corresponding module
  52 # and class names, then always index by class name.
  53 @decorators.Cache
  54 def DiscoverClasses(start_dir, top_level_dir, base_class, pattern='*',
  55                     index_by_class_name=False):
  56   """Discover all classes in |start_dir| which subclass |base_class|.
  57
  58   Base classes that contain subclasses are ignored by default.
  59
  60   Args:
  61     start_dir: The directory to recursively search.
  62     top_level_dir: The top level of the package, for importing.
  63     base_class: The base class to search for.
  64     pattern: Unix shell-style pattern for filtering the filenames to import.
  65     index_by_class_name: If True, use class name converted to
  66         lowercase_with_underscores instead of module name in return dict keys.
  67
  68   Returns:
  69     dict of {module_name: class} or {underscored_class_name: class}
  70   """
  71   modules = DiscoverModules(start_dir, top_level_dir, pattern)
  72   classes = {}
  73   for module in modules:
  74     new_classes = DiscoverClassesInModule(
  75         module, base_class, index_by_class_name)
  76     classes = dict(classes.items() + new_classes.items())
  77   return classes
  78
  79 @decorators.Cache
  80 def DiscoverClassesInModule(module, base_class, index_by_class_name=False):
  81   """Discover all classes in |module| which subclass |base_class|.
  82
  83   Base classes that contain subclasses are ignored by default.
  84
  85   Args:
  86     module: The module to search.
  87     base_class: The base class to search for.
  88     index_by_class_name: If True, use class name converted to
  89         lowercase_with_underscores instead of module name in return dict keys.
  90
  91   Returns:
  92     dict of {module_name: class} or {underscored_class_name: class}
  93   """
  94   classes = {}
  95   for _, obj in inspect.getmembers(module):
  96     # Ensure object is a class.
  97     if not inspect.isclass(obj):
  98       continue
  99     # Include only subclasses of base_class.
 100     if not issubclass(obj, base_class):
 101       continue
 102     # Exclude the base_class itself.
 103     if obj is base_class:
 104       continue
 105     # Exclude protected or private classes.
 106     if obj.__name__.startswith('_'):
 107       continue
 108     # Include only the module in which the class is defined.
 109     # If a class is imported by another module, exclude those duplicates.
 110     if obj.__module__ != module.__name__:
 111       continue
 112
 113     if index_by_class_name:
 114       key_name = camel_case.ToUnderscore(obj.__name__)
 115     else:
 116       key_name = module.__name__.split('.')[-1]
 117     classes[key_name] = obj
 118
 119   return classes
 120
 121
 122 _counter = [0]
 123 def _GetUniqueModuleName():
 124   _counter[0] += 1
 125   return "module_" + str(_counter[0])
 126
 127
 128 def IsPageSetFile(file_path):
 129   root_name, ext_name = os.path.splitext(file_path)
 130   if 'unittest' in root_name or 'page_sets/data' in root_name:
 131     return False
 132   if ext_name != '.py':
 133     return False
 134   module = util.GetPythonPageSetModule(file_path)
 135   return bool(DiscoverClassesInModule(module, page_set.PageSet))
 136
 137
 138 def GetAllPageSetFilenames(dir_path):
 139   results = []
 140   for sub_path, _, filenames in os.walk(dir_path):
 141     for f in filenames:
 142       if f.startswith('.'):
 143         continue
 144       filename = os.path.join(sub_path, f)
 145       if IsPageSetFile(filename):
 146         results.append(filename)
 147   return results