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

   1 # Copyright 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
  13
  14 @decorators.Cache
  15 def DiscoverModules(start_dir, top_level_dir, pattern='*'):
  16   """Discover all modules in |start_dir| which match |pattern|.
  17
  18   Args:
  19     start_dir: The directory to recursively search.
  20     top_level_dir: The top level of the package, for importing.
  21     pattern: Unix shell-style pattern for filtering the filenames to import.
  22
  23   Returns:
  24     list of modules.
  25   """
  26   modules = []
  27   for dir_path, _, filenames in os.walk(start_dir):
  28     for filename in filenames:
  29       # Filter out unwanted filenames.
  30       if filename.startswith('.') or filename.startswith('_'):
  31         continue
  32       if os.path.splitext(filename)[1] != '.py':
  33         continue
  34       if not fnmatch.fnmatch(filename, pattern):
  35         continue
  36
  37       # Find the module.
  38       module_rel_path = os.path.relpath(os.path.join(dir_path, filename),
  39                                         top_level_dir)
  40       module_name = re.sub(r'[/\\]', '.', os.path.splitext(module_rel_path)[0])
  41
  42       # Import the module.
  43       module = __import__(module_name, fromlist=[True])
  44
  45       modules.append(module)
  46   return modules
  47
  48
  49 # TODO(dtu): Normalize all discoverable classes to have corresponding module
  50 # and class names, then always index by class name.
  51 @decorators.Cache
  52 def DiscoverClasses(start_dir, top_level_dir, base_class, pattern='*',
  53                     index_by_class_name=False):
  54   """Discover all classes in |start_dir| which subclass |base_class|.
  55
  56   Base classes that contain subclasses are ignored by default.
  57
  58   Args:
  59     start_dir: The directory to recursively search.
  60     top_level_dir: The top level of the package, for importing.
  61     base_class: The base class to search for.
  62     pattern: Unix shell-style pattern for filtering the filenames to import.
  63     index_by_class_name: If True, use class name converted to
  64         lowercase_with_underscores instead of module name in return dict keys.
  65
  66   Returns:
  67     dict of {module_name: class} or {underscored_class_name: class}
  68   """
  69   modules = DiscoverModules(start_dir, top_level_dir, pattern)
  70   classes = {}
  71   for module in modules:
  72     new_classes = DiscoverClassesInModule(
  73         module, base_class, index_by_class_name)
  74     classes = dict(classes.items() + new_classes.items())
  75   return classes
  76
  77 @decorators.Cache
  78 def DiscoverClassesInModule(module, base_class, index_by_class_name=False):
  79   """Discover all classes in |module| which subclass |base_class|.
  80
  81   Base classes that contain subclasses are ignored by default.
  82
  83   Args:
  84     module: The module to search.
  85     base_class: The base class to search for.
  86     index_by_class_name: If True, use class name converted to
  87         lowercase_with_underscores instead of module name in return dict keys.
  88
  89   Returns:
  90     dict of {module_name: class} or {underscored_class_name: class}
  91   """
  92   classes = {}
  93   for _, obj in inspect.getmembers(module):
  94     # Ensure object is a class.
  95     if not inspect.isclass(obj):
  96       continue
  97     # Include only subclasses of base_class.
  98     if not issubclass(obj, base_class):
  99       continue
 100     # Exclude the base_class itself.
 101     if obj is base_class:
 102       continue
 103     # Exclude protected or private classes.
 104     if obj.__name__.startswith('_'):
 105       continue
 106     # Include only the module in which the class is defined.
 107     # If a class is imported by another module, exclude those duplicates.
 108     if obj.__module__ != module.__name__:
 109       continue
 110
 111     if index_by_class_name:
 112       key_name = camel_case.ToUnderscore(obj.__name__)
 113     else:
 114       key_name = module.__name__.split('.')[-1]
 115     classes[key_name] = obj
 116
 117   return classes
 118
 119
 120 _counter = [0]
 121 def _GetUniqueModuleName():
 122   _counter[0] += 1
 123   return "module_" + str(_counter[0])