1import os 2import re 3import six 4import sys 5 6# This must run on Python versions less than 2.4. 7dirname = os.path.dirname(sys.modules[__name__].__file__) 8common_dir = os.path.abspath(os.path.join(dirname, 'common_lib')) 9sys.path.insert(0, common_dir) 10import check_version 11sys.path.pop(0) 12 13 14def _get_pyversion_from_args(): 15 """Extract, format, & pop the current py_version from args, if provided.""" 16 py_version = 3 17 py_version_re = re.compile(r'--py_version=(\w+)\b') 18 19 version_found = False 20 for i, arg in enumerate(sys.argv): 21 if not arg.startswith('--py_version'): 22 continue 23 result = py_version_re.search(arg) 24 if result: 25 if version_found: 26 raise ValueError('--py_version may only be specified once.') 27 py_version = result.group(1) 28 version_found = True 29 if py_version not in ('2', '3'): 30 raise ValueError('Python version must be "2" or "3".') 31 32 # Remove the arg so other argparsers don't get grumpy. 33 sys.argv.pop(i) 34 35 return py_version 36 37 38def _desired_version(): 39 """ 40 Returns desired python version. 41 42 If the PY_VERSION env var is set, just return that. This is the case 43 when autoserv kicks of autotest on the server side via a job.run(), or 44 a process created a subprocess. 45 46 Otherwise, parse & pop the sys.argv for the '--py_version' flag. If no 47 flag is set, default to python 3. 48 49 """ 50 # Even if the arg is in the env vars, we will attempt to get it from the 51 # args, so that it can be popped prior to other argparsers hitting. 52 py_version = _get_pyversion_from_args() 53 54 if os.getenv('PY_VERSION'): 55 return int(os.getenv('PY_VERSION')) 56 57 os.environ['PY_VERSION'] = str(py_version) 58 return int(py_version) 59 60 61desired_version = _desired_version() 62if desired_version == sys.version_info.major: 63 os.environ['AUTOTEST_NO_RESTART'] = 'True' 64else: 65 # There are cases were this can be set (ie by test_that), but a subprocess 66 # is launched in the incorrect version. 67 if os.getenv('AUTOTEST_NO_RESTART'): 68 del os.environ['AUTOTEST_NO_RESTART'] 69 check_version.check_python_version(desired_version) 70 71import glob, traceback 72 73 74def import_module(module, from_where): 75 """Equivalent to 'from from_where import module' 76 Returns the corresponding module""" 77 from_module = __import__(from_where, globals(), locals(), [module]) 78 return getattr(from_module, module) 79 80 81def _autotest_logging_handle_error(self, record): 82 """Method to monkey patch into logging.Handler to replace handleError.""" 83 # The same as the default logging.Handler.handleError but also prints 84 # out the original record causing the error so there is -some- idea 85 # about which call caused the logging error. 86 import logging 87 if logging.raiseExceptions: 88 # Avoid recursion as the below output can end up back in here when 89 # something has *seriously* gone wrong in autotest. 90 logging.raiseExceptions = 0 91 sys.stderr.write('Exception occurred formatting message: ' 92 '%r using args %r\n' % (record.msg, record.args)) 93 traceback.print_stack() 94 sys.stderr.write('-' * 50 + '\n') 95 traceback.print_exc() 96 sys.stderr.write('Future logging formatting exceptions disabled.\n') 97 98 99def _monkeypatch_logging_handle_error(): 100 # Hack out logging.py* 101 logging_py = os.path.join(os.path.dirname(__file__), 'common_lib', 102 'logging.py*') 103 if glob.glob(logging_py): 104 os.system('rm -f %s' % logging_py) 105 106 # Monkey patch our own handleError into the logging module's StreamHandler. 107 # A nicer way of doing this -might- be to have our own logging module define 108 # an autotest Logger instance that added our own Handler subclass with this 109 # handleError method in it. But that would mean modifying tons of code. 110 import logging 111 assert callable(logging.Handler.handleError) 112 logging.Handler.handleError = _autotest_logging_handle_error 113 114 115def _insert_site_packages(root): 116 # Allow locally installed third party packages to be found 117 # before any that are installed on the system itself when not. 118 # running as a client. 119 # This is primarily for the benefit of frontend and tko so that they 120 # may use libraries other than those available as system packages. 121 if six.PY2: 122 sys.path.insert(0, os.path.join(root, 'site-packages')) 123 124 125import importlib 126 127ROOT_MODULE_NAME_ALLOW_LIST = ( 128 'autotest_lib', 129 'autotest_lib.client', 130) 131 132 133def _setup_top_level_symlink(base_path): 134 """Create a self pointing symlink in the base_path).""" 135 if os.path.islink(os.path.join(base_path, 'autotest_lib')): 136 return 137 os.chdir(base_path) 138 os.symlink('.', 'autotest_lib') 139 140 141def _setup_client_symlink(base_path): 142 """Setup the client symlink for the DUT. 143 144 Creates a "autotest_lib" folder in client, then creates a symlink called 145 "client" pointing back to ../, as well as an __init__ for the folder. 146 """ 147 148 def _create_client_symlink(): 149 os.chdir(autotest_lib_dir) 150 with open('__init__.py', 'w'): 151 pass 152 os.symlink('../', 'client') 153 154 autotest_lib_dir = os.path.join(base_path, 'autotest_lib') 155 link_path = os.path.join(autotest_lib_dir, 'client') 156 157 # TODO: Use os.makedirs(..., exist_ok=True) after switching to Python 3 158 if not os.path.isdir(autotest_lib_dir): 159 try: 160 os.mkdir(autotest_lib_dir) 161 except FileExistsError as e: 162 if not os.path.isdir(autotest_lib_dir): 163 raise e 164 165 if os.path.islink(link_path): 166 return 167 168 try: 169 _create_client_symlink() 170 # It's possible 2 autotest processes are running at once, and one 171 # creates the symlink in the time between checking and creating. 172 # Thus if the symlink DNE, and we cannot create it, check for its 173 # existence and exit if it exists. 174 except FileExistsError as e: 175 if os.path.islink(link_path): 176 return 177 raise e 178 179 180def _symlink_check(base_path, root_dir): 181 """Verify the required symlinks are present, and add them if not.""" 182 # Note the starting cwd to later change back to it. 183 starting_dir = os.getcwd() 184 if root_dir == 'autotest_lib': 185 _setup_top_level_symlink(base_path) 186 elif root_dir == 'autotest_lib.client': 187 _setup_client_symlink(base_path) 188 189 os.chdir(starting_dir) 190 191 192def setup(base_path, root_module_name): 193 _symlink_check(base_path, root_module_name) 194 if root_module_name not in ROOT_MODULE_NAME_ALLOW_LIST: 195 raise Exception('Unexpected root module: ' + root_module_name) 196 197 _insert_site_packages(base_path) 198 199 # Ie, server (or just not /client) 200 if root_module_name == 'autotest_lib': 201 # Base path is just x/x/x/x/autotest/files 202 _setup_autotest_lib(base_path) 203 _preimport_top_level_packages(os.path.join(base_path, 'autotest_lib'), 204 parent='autotest_lib') 205 else: # aka, in /client/ 206 if os.path.exists(os.path.join(os.path.dirname(base_path), 'server')): 207 208 # Takes you from /client/ to /files 209 # this is because on DUT there is no files/client 210 autotest_base_path = os.path.dirname(base_path) 211 212 else: 213 autotest_base_path = base_path 214 215 _setup_autotest_lib(autotest_base_path) 216 _preimport_top_level_packages(os.path.join(autotest_base_path, 217 'autotest_lib'), 218 parent='autotest_lib') 219 _preimport_top_level_packages( 220 os.path.join(autotest_base_path, 'autotest_lib', 'client'), 221 parent='autotest_lib.client', 222 ) 223 224 _monkeypatch_logging_handle_error() 225 226 227def _setup_autotest_lib(path): 228 sys.path.insert(0, path) 229 # This is a symlink back to the root directory, that does all the magic. 230 importlib.import_module('autotest_lib') 231 sys.path.pop(0) 232 233 234def _preimport_top_level_packages(root, parent): 235 # The old code to setup the packages used to fetch the top-level packages 236 # inside autotest_lib. We keep that behaviour in order to avoid having to 237 # add import statements for the top-level packages all over the codebase. 238 # 239 # e.g., 240 # import common 241 # from autotest_lib.server import utils 242 # 243 # must continue to work. The _right_ way to do that import would be. 244 # 245 # import common 246 # import autotest_lib.server 247 # from autotest_lib.server import utils 248 names = [] 249 for filename in os.listdir(root): 250 path = os.path.join(root, filename) 251 if not os.path.isdir(path): 252 continue # skip files 253 if '.' in filename: 254 continue # if "." is in the name it's not a valid package name 255 if not os.access(path, os.R_OK | os.X_OK): 256 continue # need read + exec access to make a dir importable 257 if '__init__.py' in os.listdir(path): 258 names.append(filename) 259 260 for name in names: 261 pname = parent + '.' + name 262 importlib.import_module(pname) 263 if name != 'autotest_lib': 264 sys.modules[name] = sys.modules[pname] 265