Enable re-writing of setuptools-installed plugins

Hook up the PEP 302 import hook very early in pytest startup so
that it gets installed before setuptools-installed plugins are
imported.  Also iterate over all installed plugins and mark them
for rewriting.  If an installed plugin is already imported then
a warning is issued, we can not break since that might break
existing plugins and the fallback will still be gracefull to
plain asserts.

Some existing tests are failing in this commit because of the new
warning triggered by inline pytest runs due to the hypothesis
plugin already being imported.  The tests will be fixed in the next
commit.

2 files changed, 46 insertions, 53 deletions

diff --git a/_pytest/assertion/__init__.py b/_pytest/assertion/__init__.py
index dd30e1471..746c810ee 100644
--- a/_pytest/assertion/__init__.py
+++ b/_pytest/assertion/__init__.py
@@ -5,8 +5,7 @@ import py
 import os
 import sys
 
-from _pytest.config import hookimpl
-from _pytest.monkeypatch import MonkeyPatch
+from _pytest.monkeypatch import monkeypatch
 from _pytest.assertion import util
 
 
@@ -35,10 +34,7 @@ class AssertionState:
         self.trace = config.trace.root.get("assertion")
 
 
-@hookimpl(tryfirst=True)
-def pytest_load_initial_conftests(early_config, parser, args):
-    ns, ns_unknown_args = parser.parse_known_and_unknown_args(args)
-    mode = ns.assertmode
+def install_importhook(config, mode):
     if mode == "rewrite":
         try:
             import ast  # noqa
@@ -51,37 +47,37 @@ def pytest_load_initial_conftests(early_config, parser, args):
                     sys.version_info[:3] == (2, 6, 0)):
                 mode = "reinterp"
 
-    early_config._assertstate = AssertionState(early_config, mode)
-    warn_about_missing_assertion(mode, early_config.pluginmanager)
+    config._assertstate = AssertionState(config, mode)
 
-    if mode != "plain":
-        _load_modules(mode)
-        m = MonkeyPatch()
-        early_config._cleanup.append(m.undo)
-        m.setattr(py.builtin.builtins, 'AssertionError',
-                  reinterpret.AssertionError)  # noqa
+    _load_modules(mode)
+    m = monkeypatch()
+    config._cleanup.append(m.undo)
+    m.setattr(py.builtin.builtins, 'AssertionError',
+              reinterpret.AssertionError)  # noqa
 
     hook = None
     if mode == "rewrite":
-        hook = rewrite.AssertionRewritingHook(early_config)  # noqa
+        hook = rewrite.AssertionRewritingHook(config)  # noqa
         sys.meta_path.insert(0, hook)
 
-    early_config._assertstate.hook = hook
-    early_config._assertstate.trace("configured with mode set to %r" % (mode,))
+    config._assertstate.hook = hook
+    config._assertstate.trace("configured with mode set to %r" % (mode,))
     def undo():
-        hook = early_config._assertstate.hook
+        hook = config._assertstate.hook
         if hook is not None and hook in sys.meta_path:
             sys.meta_path.remove(hook)
-    early_config.add_cleanup(undo)
+    config.add_cleanup(undo)
+    return hook
 
 
 def pytest_collection(session):
     # this hook is only called when test modules are collected
     # so for example not in the master process of pytest-xdist
     # (which does not collect test modules)
-    hook = session.config._assertstate.hook
-    if hook is not None:
-        hook.set_session(session)
+    assertstate = getattr(session.config, '_assertstate', None)
+    if assertstate:
+        if assertstate.hook is not None:
+            assertstate.hook.set_session(session)
 
 
 def _running_on_ci():
@@ -138,9 +134,10 @@ def pytest_runtest_teardown(item):
 
 
 def pytest_sessionfinish(session):
-    hook = session.config._assertstate.hook
-    if hook is not None:
-        hook.session = None
+    assertstate = getattr(session.config, '_assertstate', None)
+    if assertstate:
+        if assertstate.hook is not None:
+            assertstate.hook.set_session(None)
 
 
 def _load_modules(mode):
@@ -151,31 +148,5 @@ def _load_modules(mode):
         from _pytest.assertion import rewrite  # noqa
 
 
-def warn_about_missing_assertion(mode, pluginmanager):
-    try:
-        assert False
-    except AssertionError:
-        pass
-    else:
-        if mode == "rewrite":
-            specifically = ("assertions which are not in test modules "
-                            "will be ignored")
-        else:
-            specifically = "failing tests may report as passing"
-
-        # temporarily disable capture so we can print our warning
-        capman = pluginmanager.getplugin('capturemanager')
-        try:
-            out, err = capman.suspendcapture()
-            sys.stderr.write("WARNING: " + specifically +
-                             " because assert statements are not executed "
-                             "by the underlying Python interpreter "
-                             "(are you using python -O?)\n")
-        finally:
-            capman.resumecapture()
-            sys.stdout.write(out)
-            sys.stderr.write(err)
-
-
 # Expose this plugin's implementation for the pytest_assertrepr_compare hook
 pytest_assertrepr_compare = util.assertrepr_compare
diff --git a/_pytest/assertion/rewrite.py b/_pytest/assertion/rewrite.py
index 06944b016..50d8062ae 100644
--- a/_pytest/assertion/rewrite.py
+++ b/_pytest/assertion/rewrite.py
@@ -51,6 +51,7 @@ class AssertionRewritingHook(object):
         self.session = None
         self.modules = {}
         self._register_with_pkg_resources()
+        self._must_rewrite = set()
 
     def set_session(self, session):
         self.session = session
@@ -87,7 +88,7 @@ class AssertionRewritingHook(object):
             fn = os.path.join(pth, name.rpartition(".")[2] + ".py")
 
         fn_pypath = py.path.local(fn)
-        if not self._should_rewrite(fn_pypath, state):
+        if not self._should_rewrite(name, fn_pypath, state):
             return None
 
         # The requested module looks like a test file, so rewrite it. This is
@@ -137,7 +138,7 @@ class AssertionRewritingHook(object):
         self.modules[name] = co, pyc
         return self
 
-    def _should_rewrite(self, fn_pypath, state):
+    def _should_rewrite(self, name, fn_pypath, state):
         # always rewrite conftest files
         fn = str(fn_pypath)
         if fn_pypath.basename == 'conftest.py':
@@ -161,8 +162,29 @@ class AssertionRewritingHook(object):
                 finally:
                     self.session = session
                     del session
+        else:
+            for marked in self._must_rewrite:
+                if marked.startswith(name):
+                    return True
         return False
 
+    def mark_rewrite(self, *names):
+        """Mark import names as needing to be re-written.
+
+        The named module or package as well as any nested modules will
+        be re-written on import.
+        """
+        already_imported = set(names).intersection(set(sys.modules))
+        if already_imported:
+            self._warn_already_imported(already_imported)
+        self._must_rewrite.update(names)
+
+    def _warn_already_imported(self, names):
+        self.config.warn(
+            'P1',
+            'Modules are already imported so can not be re-written: %s' %
+            ','.join(names))
+
     def load_module(self, name):
         # If there is an existing module object named 'fullname' in
         # sys.modules, the loader must use that existing module. (Otherwise,