From c7f9fda42d475a5579aa988d8b491be3de53925b Mon Sep 17 00:00:00 2001
From: Bruno Oliveira <nicoddemus@gmail.com>
Date: Tue, 3 Dec 2019 19:49:20 -0300
Subject: [PATCH] Fix assertion rewriting module detection for egg dists

Fix #6301
---
 changelog/6301.bugfix.rst      |  1 +
 src/_pytest/config/__init__.py | 53 +++++++++++++++++++++++++++++++++-
 testing/test_config.py         | 16 ++++++----
 3 files changed, 64 insertions(+), 6 deletions(-)
 create mode 100644 changelog/6301.bugfix.rst

diff --git a/changelog/6301.bugfix.rst b/changelog/6301.bugfix.rst
new file mode 100644
index 0000000000..f13c83343f
--- /dev/null
+++ b/changelog/6301.bugfix.rst
@@ -0,0 +1 @@
+Fix assertion rewriting for egg-based distributions and ``editable`` installs (``pip install --editable``).
diff --git a/src/_pytest/config/__init__.py b/src/_pytest/config/__init__.py
index eecb29365f..038695d987 100644
--- a/src/_pytest/config/__init__.py
+++ b/src/_pytest/config/__init__.py
@@ -630,16 +630,67 @@ def __repr__(self):
 
 
 def _iter_rewritable_modules(package_files):
+    """
+    Given an iterable of file names in a source distribution, return the "names" that should
+    be marked for assertion rewrite (for example the package "pytest_mock/__init__.py" should
+    be added as "pytest_mock" in the assertion rewrite mechanism.
+
+    This function has to deal with dist-info based distributions and egg based distributions
+    (which are still very much in use for "editable" installs).
+
+    Here are the file names as seen in a dist-info based distribution:
+
+        pytest_mock/__init__.py
+        pytest_mock/_version.py
+        pytest_mock/plugin.py
+        pytest_mock.egg-info/PKG-INFO
+
+    Here are the file names as seen in an egg based distribution:
+
+        src/pytest_mock/__init__.py
+        src/pytest_mock/_version.py
+        src/pytest_mock/plugin.py
+        src/pytest_mock.egg-info/PKG-INFO
+        LICENSE
+        setup.py
+
+    We have to take in account those two distribution flavors in order to determine which
+    names should be considered for assertion rewriting.
+
+    More information:
+        https://github.com/pytest-dev/pytest-mock/issues/167
+    """
+    package_files = list(package_files)
+    seen_some = False
     for fn in package_files:
         is_simple_module = "/" not in fn and fn.endswith(".py")
         is_package = fn.count("/") == 1 and fn.endswith("__init__.py")
         if is_simple_module:
             module_name, _ = os.path.splitext(fn)
-            yield module_name
+            # we ignore "setup.py" at the root of the distribution
+            if module_name != "setup":
+                seen_some = True
+                yield module_name
         elif is_package:
             package_name = os.path.dirname(fn)
+            seen_some = True
             yield package_name
 
+    if not seen_some:
+        # at this point we did not find any packages or modules suitable for assertion
+        # rewriting, so we try again by stripping the first path component (to account for
+        # "src" based source trees for example)
+        # this approach lets us have the common case continue to be fast, as egg-distributions
+        # are rarer
+        new_package_files = []
+        for fn in package_files:
+            parts = fn.split("/")
+            new_fn = "/".join(parts[1:])
+            if new_fn:
+                new_package_files.append(new_fn)
+        if new_package_files:
+            yield from _iter_rewritable_modules(new_package_files)
+
 
 class Config:
     """
diff --git a/testing/test_config.py b/testing/test_config.py
index f146b52a45..9735fc1769 100644
--- a/testing/test_config.py
+++ b/testing/test_config.py
@@ -422,15 +422,21 @@ def test_confcutdir_check_isdir(self, testdir):
     @pytest.mark.parametrize(
         "names, expected",
         [
+            # dist-info based distributions root are files as will be put in PYTHONPATH
             (["bar.py"], ["bar"]),
-            (["foo", "bar.py"], []),
-            (["foo", "bar.pyc"], []),
-            (["foo", "__init__.py"], ["foo"]),
-            (["foo", "bar", "__init__.py"], []),
+            (["foo/bar.py"], ["bar"]),
+            (["foo/bar.pyc"], []),
+            (["foo/__init__.py"], ["foo"]),
+            (["bar/__init__.py", "xz.py"], ["bar", "xz"]),
+            (["setup.py"], []),
+            # egg based distributions root contain the files from the dist root
+            (["src/bar/__init__.py"], ["bar"]),
+            (["src/bar/__init__.py", "setup.py"], ["bar"]),
+            (["source/python/bar/__init__.py", "setup.py"], ["bar"]),
         ],
     )
     def test_iter_rewritable_modules(self, names, expected):
-        assert list(_iter_rewritable_modules(["/".join(names)])) == expected
+        assert list(_iter_rewritable_modules(names)) == expected
 
 
 class TestConfigFromdictargs: