discover.py revision a1401311d1ab56c4ed0a474bd38c108f75cb0cd9 
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
5import fnmatch
6import inspect
7import os
8import re
9
10from telemetry import decorators
11from telemetry.core import camel_case
12
13
14@decorators.Cache
15def 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
52def 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
97
98def GetAllPageSetFilenames(dir_path):
99  results = []
100  start_dir = os.path.dirname(dir_path)
101  for dirpath, _, filenames in os.walk(start_dir):
102    for f in filenames:
103      if os.path.splitext(f)[1] != '.json':
104        continue
105      filename = os.path.join(dirpath, f)
106      results.append(filename)
107  return results
108