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