fix: source modules need to be re-imported. #1232

3 files changed, 53 insertions, 29 deletions

diff --git a/coverage/inorout.py b/coverage/inorout.py
index c90e3d59..2c216ea9 100644
--- a/coverage/inorout.py
+++ b/coverage/inorout.py
@@ -18,6 +18,7 @@ from coverage.disposition import FileDisposition, disposition_init
 from coverage.exceptions import CoverageException
 from coverage.files import TreeMatcher, FnmatchMatcher, ModuleMatcher
 from coverage.files import prep_patterns, find_python_files, canonical_filename
+from coverage.misc import sys_modules_saved
 from coverage.python import source_for_file, source_for_morf
 
 
@@ -270,27 +271,28 @@ class InOrOut:
 
         # Check if the source we want to measure has been installed as a
         # third-party package.
-        for pkg in self.source_pkgs:
-            try:
-                modfile, path = file_and_path_for_module(pkg)
-                debug(f"Imported source package {pkg!r} as {modfile!r}")
-            except CoverageException as exc:
-                debug(f"Couldn't import source package {pkg!r}: {exc}")
-                continue
-            if modfile:
-                if self.third_match.match(modfile):
-                    debug(
-                        f"Source is in third-party because of source_pkg {pkg!r} at {modfile!r}"
-                    )
-                    self.source_in_third = True
-            else:
-                for pathdir in path:
-                    if self.third_match.match(pathdir):
+        with sys_modules_saved():
+            for pkg in self.source_pkgs:
+                try:
+                    modfile, path = file_and_path_for_module(pkg)
+                    debug(f"Imported source package {pkg!r} as {modfile!r}")
+                except CoverageException as exc:
+                    debug(f"Couldn't import source package {pkg!r}: {exc}")
+                    continue
+                if modfile:
+                    if self.third_match.match(modfile):
                         debug(
-                            f"Source is in third-party because of {pkg!r} path directory " +
-                            f"at {pathdir!r}"
+                            f"Source is in third-party because of source_pkg {pkg!r} at {modfile!r}"
                         )
                         self.source_in_third = True
+                else:
+                    for pathdir in path:
+                        if self.third_match.match(pathdir):
+                            debug(
+                                f"Source is in third-party because of {pkg!r} path directory " +
+                                f"at {pathdir!r}"
+                            )
+                            self.source_in_third = True
 
         for src in self.source:
             if self.third_match.match(src):
diff --git a/coverage/misc.py b/coverage/misc.py
index 9c414d88..29397537 100644
--- a/coverage/misc.py
+++ b/coverage/misc.py
@@ -3,6 +3,7 @@
 
 """Miscellaneous stuff for coverage.py."""
 
+import contextlib
 import errno
 import hashlib
 import importlib
@@ -49,6 +50,28 @@ def isolate_module(mod):
 os = isolate_module(os)
 
 
+class SysModuleSaver:
+    """Saves the contents of sys.modules, and removes new modules later."""
+    def __init__(self):
+        self.old_modules = set(sys.modules)
+
+    def restore(self):
+        """Remove any modules imported since this object started."""
+        new_modules = set(sys.modules) - self.old_modules
+        for m in new_modules:
+            del sys.modules[m]
+
+
+@contextlib.contextmanager
+def sys_modules_saved():
+    """A context manager to remove any modules imported during a block."""
+    saver = SysModuleSaver()
+    try:
+        yield
+    finally:
+        saver.restore()
+
+
 def import_third_party(modname):
     """Import a third-party module we need, but might not be installed.
 
@@ -63,16 +86,11 @@ def import_third_party(modname):
         The imported module, or None if the module couldn't be imported.
 
     """
-    try:
-        mod = importlib.import_module(modname)
-    except ImportError:
-        mod = None
-
-    imported = [m for m in sys.modules if m.startswith(modname)]
-    for name in imported:
-        del sys.modules[name]
-
-    return mod
+    with sys_modules_saved():
+        try:
+            return importlib.import_module(modname)
+        except ImportError:
+            return None
 
 
 def dummy_decorator_with_args(*args_unused, **kwargs_unused):
diff --git a/coverage/tomlconfig.py b/coverage/tomlconfig.py
index 3301acc8..4a1e322c 100644
--- a/coverage/tomlconfig.py
+++ b/coverage/tomlconfig.py
@@ -10,7 +10,11 @@ import re
 from coverage.exceptions import CoverageException
 from coverage.misc import import_third_party, substitute_variables
 
-# TOML support is an install-time extra option.
+# TOML support is an install-time extra option. (Import typing is here because
+# import_third_party will unload any module that wasn't already imported.
+# tomli imports typing, and if we unload it, later it's imported again, and on
+# Python 3.6, this causes infinite recursion.)
+import typing   # pylint: disable=unused-import, wrong-import-order
 tomli = import_third_party("tomli")