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
  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     for _, obj in inspect.getmembers(module):
  73       # Ensure object is a class.
  74       if not inspect.isclass(obj):
  75         continue
  76       # Include only subclasses of base_class.
  77       if not issubclass(obj, base_class):
  78         continue
  79       # Exclude the base_class itself.
  80       if obj is base_class:
  81         continue
  82       # Exclude protected or private classes.
  83       if obj.__name__.startswith('_'):
  84         continue
  85       # Include only the module in which the class is defined.
  86       # If a class is imported by another module, exclude those duplicates.
  87       if obj.__module__ != module.__name__:
  88         continue
  89
  90       if index_by_class_name:
  91         key_name = camel_case.ToUnderscore(obj.__name__)
  92       else:
  93         key_name = module.__name__.split('.')[-1]
  94       classes[key_name] = obj
  95
  96   return classes