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 | |
| 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. |
| 10 | """ |
| 11 | |
| 12 | from __future__ import absolute_import |
| 13 | from __future__ import division |
| 14 | from __future__ import print_function |
| 15 | |
| 16 | import imp |
| 17 | import importlib |
| 18 | import logging |
| 19 | import os |
| 20 | import site |
| 21 | import sys |
| 22 | |
| 23 | import autotest_lib |
| 24 | |
| 25 | _AUTOTEST_DIR = autotest_lib.__path__[0] |
| 26 | _SITEPKG_DIR = os.path.join(_AUTOTEST_DIR, 'site-packages') |
| 27 | |
| 28 | _setup_done = False |
| 29 | |
| 30 | logger = logging.getLogger(__name__) |
| 31 | |
| 32 | |
| 33 | def monkeypatch(): |
| 34 | """Do necessary Autotest monkeypatching. |
| 35 | |
| 36 | This should be called before any autotest_lib imports in the main |
| 37 | function in scripts. Thus, only the main function in scripts can |
| 38 | import autotest_lib. |
| 39 | |
| 40 | Library code should rely on dependency injection, falling back to |
| 41 | load(). |
| 42 | |
| 43 | This should be called no more than once. |
| 44 | |
| 45 | This adds Autotest's site-packages and modifies sys.meta_path so |
| 46 | that all common.py imports are no-ops. |
| 47 | """ |
| 48 | global _setup_done |
| 49 | assert not _setup_done |
| 50 | site.addsitedir(_SITEPKG_DIR) |
| 51 | sys.meta_path.insert(0, _CommonRemovingFinder()) |
| 52 | _setup_done = True |
| 53 | |
| 54 | |
| 55 | class _CommonRemovingFinder(object): |
| 56 | """Python import finder that neuters Autotest's common.py |
| 57 | |
| 58 | The common module is replaced with an empty module everywhere it is |
| 59 | imported. common.py should have only been imported for side |
| 60 | effects, so nothing should actually use the imported module. |
| 61 | |
| 62 | See also https://www.python.org/dev/peps/pep-0302/ |
| 63 | """ |
| 64 | |
| 65 | def find_module(self, fullname, path=None): |
| 66 | """Find module.""" |
| 67 | del path # unused |
| 68 | if not self._is_autotest_common(fullname): |
| 69 | return None |
| 70 | logger.debug('Dummying out %s import', fullname) |
| 71 | return self |
| 72 | |
| 73 | def _is_autotest_common(self, fullname): |
| 74 | return (fullname.partition('.')[0] == 'autotest_lib' |
| 75 | and fullname.rpartition('.')[-1] == 'common') |
| 76 | |
| 77 | def load_module(self, fullname): |
| 78 | """Load module.""" |
| 79 | if fullname in sys.modules: |
| 80 | return sys.modules[fullname] |
| 81 | mod = imp.new_module(fullname) |
| 82 | mod.__file__ = '<removed>' |
| 83 | mod.__loader__ = self |
| 84 | mod.__package__ = fullname.rpartition('.')[0] |
| 85 | sys.modules[fullname] = mod |
| 86 | return mod |
| 87 | |
| 88 | |
| 89 | def load(name): |
| 90 | """Import module from autotest. |
| 91 | |
| 92 | This enforces that monkeypatch() is called first. Otherwise, |
| 93 | autotest imports may or may not work. When they do work, they may |
| 94 | screw up global state. |
| 95 | """ |
| 96 | if not _setup_done: |
| 97 | raise ImportError('cannot load Autotest modules before monkeypatching') |
| 98 | relpath = name.lstrip('.') |
| 99 | return importlib.import_module('.%s' % relpath, package='autotest_lib') |