From 18652b012d8a1888917f782b6b570dcfb1389a5a Mon Sep 17 00:00:00 2001
From: Ned Batchelder <ned@nedbatchelder.com>
Date: Tue, 5 Oct 2021 19:43:00 -0400
Subject: [PATCH] fix: pretend we didn't import third-party packages we use.
 #1228

tomli couldn't use coverage themselves because we imported it early.
Cleaning sys.modules means their own imports will actually execute after
coverage has started, so their files will be properly measured.
---
 CHANGES.rst            |  5 ++++-
 coverage/misc.py       | 27 +++++++++++++++++++++++++++
 coverage/tomlconfig.py |  7 ++-----
 tests/test_misc.py     | 20 +++++++++++++++++++-
 4 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 50a81aaf3..937d2d0fa 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -22,7 +22,10 @@ This list is detailed and covers changes in each pre-release version.
 Unreleased
 ----------
 
-Nothing yet.
+- Changed an internal detail of how tomli is imported, so that tomli can use
+  coverage.py for their own test suite (`issue 1228`_).
+
+.. _issue 1228: https://github.com/nedbat/coveragepy/issues/1228
 
 
 .. _changes_60:
diff --git a/coverage/misc.py b/coverage/misc.py
index 11dad23e0..cd4a77401 100644
--- a/coverage/misc.py
+++ b/coverage/misc.py
@@ -5,6 +5,7 @@
 
 import errno
 import hashlib
+import importlib
 import importlib.util
 import inspect
 import locale
@@ -43,6 +44,32 @@ def isolate_module(mod):
 os = isolate_module(os)
 
 
+def import_third_party(modname):
+    """Import a third-party module we need, but might not be installed.
+
+    This also cleans out the module after the import, so that coverage won't
+    appear to have imported it.  This lets the third party use coverage for
+    their own tests.
+
+    Arguments:
+        modname (str): the name of the module to import.
+
+    Returns:
+        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
+
+
 def dummy_decorator_with_args(*args_unused, **kwargs_unused):
     """Dummy no-op implementation of a decorator with arguments."""
     def _decorator(func):
diff --git a/coverage/tomlconfig.py b/coverage/tomlconfig.py
index 203192c93..3301acc8e 100644
--- a/coverage/tomlconfig.py
+++ b/coverage/tomlconfig.py
@@ -8,13 +8,10 @@
 import re
 
 from coverage.exceptions import CoverageException
-from coverage.misc import substitute_variables
+from coverage.misc import import_third_party, substitute_variables
 
 # TOML support is an install-time extra option.
-try:
-    import tomli
-except ImportError:         # pragma: not covered
-    tomli = None
+tomli = import_third_party("tomli")
 
 
 class TomlDecodeError(Exception):
diff --git a/tests/test_misc.py b/tests/test_misc.py
index 3858c4f8b..077c24344 100644
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -3,11 +3,13 @@
 
 """Tests of miscellaneous stuff."""
 
+import sys
+
 import pytest
 
 from coverage.exceptions import CoverageException
 from coverage.misc import contract, dummy_decorator_with_args, file_be_gone
-from coverage.misc import Hasher, one_of, substitute_variables
+from coverage.misc import Hasher, one_of, substitute_variables, import_third_party
 from coverage.misc import USE_CONTRACTS
 
 from tests.coveragetest import CoverageTest
@@ -155,3 +157,19 @@ def test_substitute_variables_errors(text):
         substitute_variables(text, VARS)
     assert text in str(exc_info.value)
     assert "Variable NOTHING is undefined" in str(exc_info.value)
+
+
+class ImportThirdPartyTest(CoverageTest):
+    """Test import_third_party."""
+
+    run_in_temp_dir = False
+
+    def test_success(self):
+        mod = import_third_party("pytest")
+        assert mod.__name__ == "pytest"
+        assert "pytest" not in sys.modules
+
+    def test_failure(self):
+        mod = import_third_party("xyzzy")
+        assert mod is None
+        assert "xyzzy" not in sys.modules