From d53e449296e214ab09632c4305f4ea3368d0b303 Mon Sep 17 00:00:00 2001 From: Fabio Zadrozny Date: Fri, 31 Aug 2018 12:27:08 -0300 Subject: [PATCH] Improve performance of assertion rewriting. Fixes #3918 --- .gitignore | 3 ++ AUTHORS | 1 + changelog/3918.bugfix.rst | 1 + src/_pytest/assertion/rewrite.py | 70 +++++++++++++++++++++++++------- src/_pytest/main.py | 5 ++- 5 files changed, 63 insertions(+), 17 deletions(-) create mode 100644 changelog/3918.bugfix.rst diff --git a/.gitignore b/.gitignore index afb6bf9fd..b5465e2e5 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,6 @@ env/ .ropeproject .idea .hypothesis +.pydevproject +.project +.settings \ No newline at end of file diff --git a/AUTHORS b/AUTHORS index 1641ea15e..385310e8d 100644 --- a/AUTHORS +++ b/AUTHORS @@ -71,6 +71,7 @@ Endre Galaczi Eric Hunsberger Eric Siegerman Erik M. Bray +Fabio Zadrozny Feng Ma Florian Bruhin Floris Bruynooghe diff --git a/changelog/3918.bugfix.rst b/changelog/3918.bugfix.rst new file mode 100644 index 000000000..95a90a000 --- /dev/null +++ b/changelog/3918.bugfix.rst @@ -0,0 +1 @@ +Improve performance of assertion rewriting. \ No newline at end of file diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py index a48a931ac..706fa2120 100644 --- a/src/_pytest/assertion/rewrite.py +++ b/src/_pytest/assertion/rewrite.py @@ -67,14 +67,20 @@ class AssertionRewritingHook(object): # flag to guard against trying to rewrite a pyc file while we are already writing another pyc file, # which might result in infinite recursion (#3506) self._writing_pyc = False + self._basenames_to_check_rewrite = set('conftest',) + self._marked_for_rewrite_cache = {} + self._session_paths_checked = False def set_session(self, session): self.session = session + self._session_paths_checked = False def find_module(self, name, path=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] @@ -166,6 +172,41 @@ class AssertionRewritingHook(object): self.modules[name] = co, pyc return self + 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 + 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. + """ + if not self._session_paths_checked and self.session is not None \ + and hasattr(self.session, '_initialpaths'): + self._session_paths_checked = True + for path in self.session._initialpaths: + # Make something as c:/projects/my_project/path.py -> + # ['c:', 'projects', 'my_project', 'path.py'] + parts = str(path).split(os.path.sep) + # add 'path' to basenames to be checked. + self._basenames_to_check_rewrite.add(os.path.splitext(parts[-1])[0]) + + # Note: conftest already by default in _basenames_to_check_rewrite. + parts = name.split('.') + if parts[-1] in self._basenames_to_check_rewrite: + return False + + # For matching the name it must be as if it was a filename. + parts[-1] = parts[-1] + '.py' + fn_pypath = py.path.local(os.path.sep.join(parts)) + for pat in self.fnpats: + if fn_pypath.fnmatch(pat): + return False + + if self._is_marked_for_rewrite(name, state): + return False + + state.trace("early skip of rewriting module: %s" % (name,)) + return True + def _should_rewrite(self, name, fn_pypath, state): # always rewrite conftest files fn = str(fn_pypath) @@ -185,12 +226,20 @@ class AssertionRewritingHook(object): state.trace("matched test file %r" % (fn,)) return True - for marked in self._must_rewrite: - if name == marked or name.startswith(marked + "."): - state.trace("matched marked file %r (from %r)" % (name, marked)) - return True + return self._is_marked_for_rewrite(name, state) - return False + def _is_marked_for_rewrite(self, name, state): + try: + return self._marked_for_rewrite_cache[name] + except KeyError: + for marked in self._must_rewrite: + if name == marked or name.startswith(marked + "."): + state.trace("matched marked file %r (from %r)" % (name, marked)) + self._marked_for_rewrite_cache[name] = True + return True + + self._marked_for_rewrite_cache[name] = False + return False def mark_rewrite(self, *names): """Mark import names as needing to be rewritten. @@ -207,6 +256,7 @@ class AssertionRewritingHook(object): ): self._warn_already_imported(name) self._must_rewrite.update(names) + self._marked_for_rewrite_cache.clear() def _warn_already_imported(self, name): self.config.warn( @@ -239,16 +289,6 @@ class AssertionRewritingHook(object): raise return sys.modules[name] - def is_package(self, name): - try: - fd, fn, desc = imp.find_module(name) - except ImportError: - return False - if fd is not None: - fd.close() - tp = desc[2] - return tp == imp.PKG_DIRECTORY - @classmethod def _register_with_pkg_resources(cls): """ diff --git a/src/_pytest/main.py b/src/_pytest/main.py index 947c6aa4b..c24ac8514 100644 --- a/src/_pytest/main.py +++ b/src/_pytest/main.py @@ -441,13 +441,14 @@ class Session(nodes.FSCollector): self.trace("perform_collect", self, args) self.trace.root.indent += 1 self._notfound = [] - self._initialpaths = set() + initialpaths = [] self._initialparts = [] self.items = items = [] for arg in args: parts = self._parsearg(arg) self._initialparts.append(parts) - self._initialpaths.add(parts[0]) + initialpaths.append(parts[0]) + self._initialpaths = frozenset(initialpaths) rep = collect_one_node(self) self.ihook.pytest_collectreport(report=rep) self.trace.root.indent -= 1