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.core import camel_case
  11
  12
  13 def DiscoverModules(start_dir, top_level_dir, pattern='*'):
  14   """Discover all modules in |start_dir| which match |pattern|.
  15
  16   Args:
  17     start_dir: The directory to recursively search.
  18     top_level_dir: The top level of the package, for importing.
  19     pattern: Unix shell-style pattern for filtering the filenames to import.
  20
  21   Returns:
  22     list of modules.
  23   """
  24   modules = []
  25   for dir_path, _, filenames in os.walk(start_dir):
  26     for filename in filenames:
  27       # Filter out unwanted filenames.
  28       if filename.startswith('.') or filename.startswith('_'):
  29         continue
  30       if os.path.splitext(filename)[1] != '.py':
  31         continue
  32       if not fnmatch.fnmatch(filename, pattern):
  33         continue
  34
  35       # Find the module.
  36       module_rel_path = os.path.relpath(os.path.join(dir_path, filename),
  37                                         top_level_dir)
  38       module_name = re.sub(r'[/\\]', '.', os.path.splitext(module_rel_path)[0])
  39
  40       # Import the module.
  41       module = __import__(module_name, fromlist=[True])
  42
  43       modules.append(module)
  44   return modules
  45
  46
  47 # TODO(dtu): Normalize all discoverable classes to have corresponding module
  48 # and class names, then always index by class name.
  49 def DiscoverClasses(start_dir, top_level_dir, base_class, pattern='*',
  50                     index_by_class_name=False):
  51   """Discover all classes in |start_dir| which subclass |base_class|.
  52
  53   Base classes that contain subclasses are ignored by default.
  54
  55   Args:
  56     start_dir: The directory to recursively search.
  57     top_level_dir: The top level of the package, for importing.
  58     base_class: The base class to search for.
  59     pattern: Unix shell-style pattern for filtering the filenames to import.
  60     index_by_class_name: If True, use class name converted to
  61         lowercase_with_underscores instead of module name in return dict keys.
  62
  63   Returns:
  64     dict of {module_name: class} or {underscored_class_name: class}
  65   """
  66   modules = DiscoverModules(start_dir, top_level_dir, pattern)
  67   classes = {}
  68   for module in modules:
  69     for _, obj in inspect.getmembers(module):
  70       if (inspect.isclass(obj) and obj is not base_class and
  71           issubclass(obj, base_class) and obj.__module__ == module.__name__
  72           and len(obj.__subclasses__()) == 0):
  73         if index_by_class_name:
  74           key_name = camel_case.ToUnderscore(obj.__name__)
  75         else:
  76           key_name = module.__name__.split('.')[-1]
  77         classes[key_name] = obj
  78
  79   return classes