1__author__ = "jadmanski@google.com (John Admanski)"
2
3import os, sys
4
5# This must run on Python versions less than 2.4.
6dirname = os.path.dirname(sys.modules[__name__].__file__)
7common_dir = os.path.abspath(os.path.join(dirname, "common_lib"))
8sys.path.insert(0, common_dir)
9import check_version
10sys.path.pop(0)
11check_version.check_python_version()
12
13import new, glob, traceback
14
15
16def _create_module(name):
17    """Create a single top-level module"""
18    module = new.module(name)
19    sys.modules[name] = module
20    return module
21
22
23def _create_module_and_parents(name):
24    """Create a module, and all the necessary parents"""
25    parts = name.split(".")
26    # first create the top-level module
27    parent = _create_module(parts[0])
28    created_parts = [parts[0]]
29    parts.pop(0)
30    # now, create any remaining child modules
31    while parts:
32        child_name = parts.pop(0)
33        module = new.module(child_name)
34        setattr(parent, child_name, module)
35        created_parts.append(child_name)
36        sys.modules[".".join(created_parts)] = module
37        parent = module
38
39
40def _import_children_into_module(parent_module_name, path):
41    """Import all the packages on a path into a parent module"""
42    # find all the packages at 'path'
43    names = []
44    for filename in os.listdir(path):
45        full_name = os.path.join(path, filename)
46        if not os.path.isdir(full_name):
47            continue   # skip files
48        if "." in filename:
49            continue   # if "." is in the name it's not a valid package name
50        if not os.access(full_name, os.R_OK | os.X_OK):
51            continue   # need read + exec access to make a dir importable
52        if "__init__.py" in os.listdir(full_name):
53            names.append(filename)
54    # import all the packages and insert them into 'parent_module'
55    sys.path.insert(0, path)
56    for name in names:
57        module = __import__(name)
58        # add the package to the parent
59        parent_module = sys.modules[parent_module_name]
60        setattr(parent_module, name, module)
61        full_name = parent_module_name + "." + name
62        sys.modules[full_name] = module
63    # restore the system path
64    sys.path.pop(0)
65
66
67def import_module(module, from_where):
68    """Equivalent to 'from from_where import module'
69    Returns the corresponding module"""
70    from_module = __import__(from_where, globals(), locals(), [module])
71    return getattr(from_module, module)
72
73
74def _autotest_logging_handle_error(self, record):
75    """Method to monkey patch into logging.Handler to replace handleError."""
76    # The same as the default logging.Handler.handleError but also prints
77    # out the original record causing the error so there is -some- idea
78    # about which call caused the logging error.
79    import logging
80    if logging.raiseExceptions:
81        # Avoid recursion as the below output can end up back in here when
82        # something has *seriously* gone wrong in autotest.
83        logging.raiseExceptions = 0
84        sys.stderr.write('Exception occurred formatting message: '
85                         '%r using args %r\n' % (record.msg, record.args))
86        traceback.print_stack()
87        sys.stderr.write('-' * 50 + '\n')
88        traceback.print_exc()
89        sys.stderr.write('Future logging formatting exceptions disabled.\n')
90
91
92def _monkeypatch_logging_handle_error():
93    # Hack out logging.py*
94    logging_py = os.path.join(os.path.dirname(__file__), "common_lib",
95                              "logging.py*")
96    if glob.glob(logging_py):
97        os.system("rm -f %s" % logging_py)
98
99    # Monkey patch our own handleError into the logging module's StreamHandler.
100    # A nicer way of doing this -might- be to have our own logging module define
101    # an autotest Logger instance that added our own Handler subclass with this
102    # handleError method in it.  But that would mean modifying tons of code.
103    import logging
104    assert callable(logging.Handler.handleError)
105    logging.Handler.handleError = _autotest_logging_handle_error
106
107
108def setup(base_path, root_module_name=""):
109    """
110    Perform all the necessary setup so that all the packages at
111    'base_path' can be imported via "import root_module_name.package".
112    If root_module_name is empty, then all the packages at base_path
113    are inserted as top-level packages.
114
115    Also, setup all the common.* aliases for modules in the common
116    library.
117
118    The setup must be different if you are running on an Autotest server
119    or on a test machine that just has the client directories installed.
120    """
121    # Hack... Any better ideas?
122    if (root_module_name == 'autotest_lib.client' and
123        os.path.exists(os.path.join(os.path.dirname(__file__),
124                                    '..', 'server'))):
125        root_module_name = 'autotest_lib'
126        base_path = os.path.abspath(os.path.join(base_path, '..'))
127
128    _create_module_and_parents(root_module_name)
129    _import_children_into_module(root_module_name, base_path)
130
131    if root_module_name == 'autotest_lib':
132        # Allow locally installed third party packages to be found
133        # before any that are installed on the system itself when not.
134        # running as a client.
135        # This is primarily for the benefit of frontend and tko so that they
136        # may use libraries other than those available as system packages.
137        sys.path.insert(0, os.path.join(base_path, "site-packages"))
138
139    _monkeypatch_logging_handle_error()
140