1# Copyright 2017 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"""Kludges to support legacy Autotest code. 6 7Autotest imports should be done by calling monkeypatch() first and then 8calling load(). monkeypatch() should only be called once from a 9script's main function. 10 11chromite imports should be done with chromite_load(), and any third 12party packages should be imported with deps_load(). The reason for this 13is to present a clear API for these unsafe imports, making it easier to 14identify which imports are currently unsafe. Eventually, everything 15should be moved to virtualenv, but that will not be in the near future. 16""" 17 18from __future__ import absolute_import 19from __future__ import division 20from __future__ import print_function 21 22import ast 23import contextlib 24import imp 25import importlib 26import logging 27import os 28import site 29import subprocess 30import sys 31 32import autotest_lib 33 34AUTOTEST_DIR = autotest_lib.__path__[0] 35_SITEPKG_DIR = os.path.join(AUTOTEST_DIR, 'site-packages') 36_SYSTEM_PYTHON = '/usr/bin/python2.7' 37 38_setup_done = False 39 40logger = logging.getLogger(__name__) 41 42 43def monkeypatch(): 44 """Monkeypatch everything needed to import Autotest. 45 46 This should be called before any calls to load(). Only the main 47 function in scripts should call this function. 48 """ 49 with _global_setup(): 50 _monkeypatch_body() 51 52 53@contextlib.contextmanager 54def _global_setup(): 55 """Context manager for checking and setting global _setup_done variable.""" 56 global _setup_done 57 assert not _setup_done 58 try: 59 yield 60 except Exception: # pragma: no cover 61 # We cannot recover from this since we leave the interpreter in 62 # an unknown state. 63 logger.exception('Uncaught exception escaped Autotest setup') 64 sys.exit(1) 65 else: 66 _setup_done = True 67 68 69def _monkeypatch_body(): 70 """The body of monkeypatch() running within _global_setup() context.""" 71 # Add Autotest's site-packages. 72 site.addsitedir(_SITEPKG_DIR) 73 74 # Dummy out common imports as they may cause problems. 75 sys.meta_path.insert(0, _CommonRemovingFinder()) 76 77 # Add chromite's third-party to the import path (chromite does this 78 # on import). 79 try: 80 importlib.import_module('chromite') 81 except ImportError: 82 # Moblab does not run build_externals; dependencies like 83 # chromite are installed system-wide. 84 logger.info("""\ 85Could not find chromite; adding system packages and retrying \ 86(This should only happen on Moblab)""") 87 for d in _system_site_packages(): 88 site.addsitedir(d) 89 importlib.import_module('chromite') 90 91 # Set up Django environment variables. 92 importlib.import_module('autotest_lib.frontend.setup_django_environment') 93 94 # Make Django app paths absolute. 95 settings = importlib.import_module('autotest_lib.frontend.settings') 96 settings.INSTALLED_APPS = ( 97 'autotest_lib.frontend.afe', 98 'autotest_lib.frontend.tko', 99 'django.contrib.admin', 100 'django.contrib.auth', 101 'django.contrib.contenttypes', 102 'django.contrib.sessions', 103 'django.contrib.sites', 104 ) 105 106 # drone_utility uses this. 107 common = importlib.import_module('autotest_lib.scheduler.common') 108 common.autotest_dir = AUTOTEST_DIR 109 110 111def _system_site_packages(): 112 """Get list of system site-package directories. 113 114 This is needed for Moblab because dependencies are installed 115 system-wide instead of using build_externals.py. 116 """ 117 output = subprocess.check_output([ 118 _SYSTEM_PYTHON, '-c', 119 'import site; print repr(site.getsitepackages())']) 120 return ast.literal_eval(output) 121 122 123class _CommonRemovingFinder(object): 124 """Python import finder that neuters Autotest's common.py 125 126 The common module is replaced with an empty module everywhere it is 127 imported. common.py should have only been imported for side 128 effects, so nothing should actually use the imported module. 129 130 See also https://www.python.org/dev/peps/pep-0302/ 131 """ 132 133 def find_module(self, fullname, path=None): 134 """Find module.""" 135 del path # unused 136 if not self._is_autotest_common(fullname): 137 return None 138 logger.debug('Dummying out %s import', fullname) 139 return self 140 141 def _is_autotest_common(self, fullname): 142 return (fullname.partition('.')[0] == 'autotest_lib' 143 and fullname.rpartition('.')[-1] == 'common') 144 145 def load_module(self, fullname): 146 """Load module.""" 147 if fullname in sys.modules: # pragma: no cover 148 return sys.modules[fullname] 149 mod = imp.new_module(fullname) 150 mod.__file__ = '<removed>' 151 mod.__loader__ = self 152 mod.__package__ = fullname.rpartition('.')[0] 153 sys.modules[fullname] = mod 154 return mod 155 156 157def load(name): 158 """Import module from autotest. 159 160 This enforces that monkeypatch() is called first. 161 162 @param name: name of module as string, e.g., 'frontend.afe.models' 163 """ 164 return _load('autotest_lib.%s' % name) 165 166 167def chromite_load(name): 168 """Import module from chromite.lib. 169 170 This enforces that monkeypatch() is called first. 171 172 @param name: name of module as string, e.g., 'metrics' 173 """ 174 return _load('chromite.lib.%s' % name) 175 176 177def deps_load(name): 178 """Import module from dependencies, e.g. site-package. 179 180 This enforces that monkeypatch() is called first. 181 182 @param name: name of module as string, e.g., 'metrics' 183 """ 184 assert not name.startswith('autotest_lib') 185 assert not name.startswith('chromite.lib') 186 return _load(name) 187 188 189def _load(name): 190 """Import a module. 191 192 This enforces that monkeypatch() is called first. 193 194 @param name: name of module as string 195 """ 196 if not _setup_done: 197 raise ImportError('cannot load chromite modules before monkeypatching') 198 return importlib.import_module(name) 199