Switch from deprecated imp to importlib

This commit is contained in:
Anthony Sottile
2019-06-22 15:02:32 -07:00
parent 3d01dd3adf
commit 4cd08f9b52
13 changed files with 182 additions and 242 deletions
+79 -160
View File
@@ -1,18 +1,16 @@
"""Rewrite assertion AST to produce nice error messages"""
import ast
import errno
import imp
import importlib.machinery
import importlib.util
import itertools
import marshal
import os
import re
import struct
import sys
import types
from importlib.util import spec_from_file_location
import atomicwrites
import py
from _pytest._io.saferepr import saferepr
from _pytest.assertion import util
@@ -23,23 +21,13 @@ from _pytest.pathlib import fnmatch_ex
from _pytest.pathlib import PurePath
# pytest caches rewritten pycs in __pycache__.
if hasattr(imp, "get_tag"):
PYTEST_TAG = imp.get_tag() + "-PYTEST"
else:
if hasattr(sys, "pypy_version_info"):
impl = "pypy"
else:
impl = "cpython"
ver = sys.version_info
PYTEST_TAG = "{}-{}{}-PYTEST".format(impl, ver[0], ver[1])
del ver, impl
PYTEST_TAG = "{}-PYTEST".format(sys.implementation.cache_tag)
PYC_EXT = ".py" + (__debug__ and "c" or "o")
PYC_TAIL = "." + PYTEST_TAG + PYC_EXT
class AssertionRewritingHook:
"""PEP302 Import hook which rewrites asserts."""
"""PEP302/PEP451 import hook which rewrites asserts."""
def __init__(self, config):
self.config = config
@@ -48,7 +36,6 @@ class AssertionRewritingHook:
except ValueError:
self.fnpats = ["test_*.py", "*_test.py"]
self.session = None
self.modules = {}
self._rewritten_names = set()
self._must_rewrite = set()
# flag to guard against trying to rewrite a pyc file while we are already writing another pyc file,
@@ -62,55 +49,51 @@ class AssertionRewritingHook:
self.session = session
self._session_paths_checked = False
def _imp_find_module(self, name, path=None):
"""Indirection so we can mock calls to find_module originated from the hook during testing"""
return imp.find_module(name, path)
# Indirection so we can mock calls to find_spec originated from the hook during testing
_find_spec = importlib.machinery.PathFinder.find_spec
def find_module(self, name, path=None):
def find_spec(self, name, path=None, target=None):
if self._writing_pyc:
return None
state = self.config._assertstate
if self._early_rewrite_bailout(name, state):
return None
state.trace("find_module called for: %s" % name)
names = name.rsplit(".", 1)
lastname = names[-1]
pth = None
if path is not None:
# Starting with Python 3.3, path is a _NamespacePath(), which
# causes problems if not converted to list.
path = list(path)
if len(path) == 1:
pth = path[0]
if pth is None:
try:
fd, fn, desc = self._imp_find_module(lastname, path)
except ImportError:
return None
if fd is not None:
fd.close()
tp = desc[2]
if tp == imp.PY_COMPILED:
if hasattr(imp, "source_from_cache"):
try:
fn = imp.source_from_cache(fn)
except ValueError:
# Python 3 doesn't like orphaned but still-importable
# .pyc files.
fn = fn[:-1]
else:
fn = fn[:-1]
elif tp != imp.PY_SOURCE:
# Don't know what this is.
return None
else:
fn = os.path.join(pth, name.rpartition(".")[2] + ".py")
fn_pypath = py.path.local(fn)
if not self._should_rewrite(name, fn_pypath, state):
spec = self._find_spec(name, path)
if (
# the import machinery could not find a file to import
spec is None
# this is a namespace package (without `__init__.py`)
# there's nothing to rewrite there
# python3.5 - python3.6: `namespace`
# python3.7+: `None`
or spec.origin in {None, "namespace"}
# if the file doesn't exist, we can't rewrite it
or not os.path.exists(spec.origin)
):
return None
else:
fn = spec.origin
if not self._should_rewrite(name, fn, state):
return None
self._rewritten_names.add(name)
return importlib.util.spec_from_file_location(
name,
fn,
loader=self,
submodule_search_locations=spec.submodule_search_locations,
)
def create_module(self, spec):
return None # default behaviour is fine
def exec_module(self, module):
fn = module.__spec__.origin
state = self.config._assertstate
self._rewritten_names.add(module.__name__)
# The requested module looks like a test file, so rewrite it. This is
# the most magical part of the process: load the source, rewrite the
@@ -121,7 +104,7 @@ class AssertionRewritingHook:
# cached pyc is always a complete, valid pyc. Operations on it must be
# atomic. POSIX's atomic rename comes in handy.
write = not sys.dont_write_bytecode
cache_dir = os.path.join(fn_pypath.dirname, "__pycache__")
cache_dir = os.path.join(os.path.dirname(fn), "__pycache__")
if write:
try:
os.mkdir(cache_dir)
@@ -132,26 +115,23 @@ class AssertionRewritingHook:
# common case) or it's blocked by a non-dir node. In the
# latter case, we'll ignore it in _write_pyc.
pass
elif e in [errno.ENOENT, errno.ENOTDIR]:
elif e in {errno.ENOENT, errno.ENOTDIR}:
# One of the path components was not a directory, likely
# because we're in a zip file.
write = False
elif e in [errno.EACCES, errno.EROFS, errno.EPERM]:
state.trace("read only directory: %r" % fn_pypath.dirname)
elif e in {errno.EACCES, errno.EROFS, errno.EPERM}:
state.trace("read only directory: %r" % os.path.dirname(fn))
write = False
else:
raise
cache_name = fn_pypath.basename[:-3] + PYC_TAIL
cache_name = os.path.basename(fn)[:-3] + PYC_TAIL
pyc = os.path.join(cache_dir, cache_name)
# Notice that even if we're in a read-only directory, I'm going
# to check for a cached pyc. This may not be optimal...
co = _read_pyc(fn_pypath, pyc, state.trace)
co = _read_pyc(fn, pyc, state.trace)
if co is None:
state.trace("rewriting {!r}".format(fn))
source_stat, co = _rewrite_test(self.config, fn_pypath)
if co is None:
# Probably a SyntaxError in the test.
return None
source_stat, co = _rewrite_test(fn)
if write:
self._writing_pyc = True
try:
@@ -160,13 +140,11 @@ class AssertionRewritingHook:
self._writing_pyc = False
else:
state.trace("found cached rewritten pyc for {!r}".format(fn))
self.modules[name] = co, pyc
return self
exec(co, module.__dict__)
def _early_rewrite_bailout(self, name, state):
"""
This is a fast way to get out of rewriting modules. Profiling has
shown that the call to imp.find_module (inside of the find_module
"""This is a fast way to get out of rewriting modules. Profiling has
shown that the call to PathFinder.find_spec (inside of the find_spec
from this class) is a major slowdown, so, this method tries to
filter what we're sure won't be rewritten before getting to it.
"""
@@ -201,10 +179,9 @@ class AssertionRewritingHook:
state.trace("early skip of rewriting module: {}".format(name))
return True
def _should_rewrite(self, name, fn_pypath, state):
def _should_rewrite(self, name, fn, state):
# always rewrite conftest files
fn = str(fn_pypath)
if fn_pypath.basename == "conftest.py":
if os.path.basename(fn) == "conftest.py":
state.trace("rewriting conftest file: {!r}".format(fn))
return True
@@ -217,8 +194,9 @@ class AssertionRewritingHook:
# modules not passed explicitly on the command line are only
# rewritten if they match the naming convention for test files
fn_path = PurePath(fn)
for pat in self.fnpats:
if fn_pypath.fnmatch(pat):
if fnmatch_ex(pat, fn_path):
state.trace("matched test file {!r}".format(fn))
return True
@@ -249,9 +227,10 @@ class AssertionRewritingHook:
set(names).intersection(sys.modules).difference(self._rewritten_names)
)
for name in already_imported:
mod = sys.modules[name]
if not AssertionRewriter.is_rewrite_disabled(
sys.modules[name].__doc__ or ""
):
mod.__doc__ or ""
) and not isinstance(mod.__loader__, type(self)):
self._warn_already_imported(name)
self._must_rewrite.update(names)
self._marked_for_rewrite_cache.clear()
@@ -268,45 +247,8 @@ class AssertionRewritingHook:
stacklevel=5,
)
def load_module(self, name):
co, pyc = self.modules.pop(name)
if name in sys.modules:
# If there is an existing module object named 'fullname' in
# sys.modules, the loader must use that existing module. (Otherwise,
# the reload() builtin will not work correctly.)
mod = sys.modules[name]
else:
# I wish I could just call imp.load_compiled here, but __file__ has to
# be set properly. In Python 3.2+, this all would be handled correctly
# by load_compiled.
mod = sys.modules[name] = imp.new_module(name)
try:
mod.__file__ = co.co_filename
# Normally, this attribute is 3.2+.
mod.__cached__ = pyc
mod.__loader__ = self
# Normally, this attribute is 3.4+
mod.__spec__ = spec_from_file_location(name, co.co_filename, loader=self)
exec(co, mod.__dict__)
except: # noqa
if name in sys.modules:
del sys.modules[name]
raise
return sys.modules[name]
def is_package(self, name):
try:
fd, fn, desc = self._imp_find_module(name)
except ImportError:
return False
if fd is not None:
fd.close()
tp = desc[2]
return tp == imp.PKG_DIRECTORY
def get_data(self, pathname):
"""Optional PEP302 get_data API.
"""
"""Optional PEP302 get_data API."""
with open(pathname, "rb") as f:
return f.read()
@@ -314,15 +256,13 @@ class AssertionRewritingHook:
def _write_pyc(state, co, source_stat, pyc):
# Technically, we don't have to have the same pyc format as
# (C)Python, since these "pycs" should never be seen by builtin
# import. However, there's little reason deviate, and I hope
# sometime to be able to use imp.load_compiled to load them. (See
# the comment in load_module above.)
# import. However, there's little reason deviate.
try:
with atomicwrites.atomic_write(pyc, mode="wb", overwrite=True) as fp:
fp.write(imp.get_magic())
fp.write(importlib.util.MAGIC_NUMBER)
# as of now, bytecode header expects 32-bit numbers for size and mtime (#4903)
mtime = int(source_stat.mtime) & 0xFFFFFFFF
size = source_stat.size & 0xFFFFFFFF
mtime = int(source_stat.st_mtime) & 0xFFFFFFFF
size = source_stat.st_size & 0xFFFFFFFF
# "<LL" stands for 2 unsigned longs, little-ending
fp.write(struct.pack("<LL", mtime, size))
fp.write(marshal.dumps(co))
@@ -335,35 +275,14 @@ def _write_pyc(state, co, source_stat, pyc):
return True
RN = b"\r\n"
N = b"\n"
cookie_re = re.compile(r"^[ \t\f]*#.*coding[:=][ \t]*[-\w.]+")
BOM_UTF8 = "\xef\xbb\xbf"
def _rewrite_test(config, fn):
"""Try to read and rewrite *fn* and return the code object."""
state = config._assertstate
try:
stat = fn.stat()
source = fn.read("rb")
except EnvironmentError:
return None, None
try:
tree = ast.parse(source, filename=fn.strpath)
except SyntaxError:
# Let this pop up again in the real import.
state.trace("failed to parse: {!r}".format(fn))
return None, None
rewrite_asserts(tree, fn, config)
try:
co = compile(tree, fn.strpath, "exec", dont_inherit=True)
except SyntaxError:
# It's possible that this error is from some bug in the
# assertion rewriting, but I don't know of a fast way to tell.
state.trace("failed to compile: {!r}".format(fn))
return None, None
def _rewrite_test(fn):
"""read and rewrite *fn* and return the code object."""
stat = os.stat(fn)
with open(fn, "rb") as f:
source = f.read()
tree = ast.parse(source, filename=fn)
rewrite_asserts(tree, fn)
co = compile(tree, fn, "exec", dont_inherit=True)
return stat, co
@@ -378,8 +297,9 @@ def _read_pyc(source, pyc, trace=lambda x: None):
return None
with fp:
try:
mtime = int(source.mtime())
size = source.size()
stat_result = os.stat(source)
mtime = int(stat_result.st_mtime)
size = stat_result.st_size
data = fp.read(12)
except EnvironmentError as e:
trace("_read_pyc({}): EnvironmentError {}".format(source, e))
@@ -387,7 +307,7 @@ def _read_pyc(source, pyc, trace=lambda x: None):
# Check for invalid or out of date pyc file.
if (
len(data) != 12
or data[:4] != imp.get_magic()
or data[:4] != importlib.util.MAGIC_NUMBER
or struct.unpack("<LL", data[4:]) != (mtime & 0xFFFFFFFF, size & 0xFFFFFFFF)
):
trace("_read_pyc(%s): invalid or out of date pyc" % source)
@@ -403,9 +323,9 @@ def _read_pyc(source, pyc, trace=lambda x: None):
return co
def rewrite_asserts(mod, module_path=None, config=None):
def rewrite_asserts(mod, module_path=None):
"""Rewrite the assert statements in mod."""
AssertionRewriter(module_path, config).run(mod)
AssertionRewriter(module_path).run(mod)
def _saferepr(obj):
@@ -586,10 +506,9 @@ class AssertionRewriter(ast.NodeVisitor):
"""
def __init__(self, module_path, config):
def __init__(self, module_path):
super().__init__()
self.module_path = module_path
self.config = config
def run(self, mod):
"""Find all assert statements in *mod* and rewrite them."""
@@ -758,7 +677,7 @@ class AssertionRewriter(ast.NodeVisitor):
"assertion is always true, perhaps remove parentheses?"
),
category=None,
filename=str(self.module_path),
filename=self.module_path,
lineno=assert_.lineno,
)
@@ -817,7 +736,7 @@ class AssertionRewriter(ast.NodeVisitor):
AST_NONE = ast.parse("None").body[0].value
val_is_none = ast.Compare(node, [ast.Is()], [AST_NONE])
send_warning = ast.parse(
"""
"""\
from _pytest.warning_types import PytestAssertRewriteWarning
from warnings import warn_explicit
warn_explicit(
@@ -827,7 +746,7 @@ warn_explicit(
lineno={lineno},
)
""".format(
filename=module_path.strpath, lineno=lineno
filename=module_path, lineno=lineno
)
).body
return ast.If(val_is_none, send_warning, [])