@@ -622,16 +622,68 @@ def __repr__(self):
622622
623623
624624def _iter_rewritable_modules (package_files ):
625+ """
626+ Given an iterable of file names in a source distribution, return the "names" that should
627+ be marked for assertion rewrite (for example the package "pytest_mock/__init__.py" should
628+ be added as "pytest_mock" in the assertion rewrite mechanism.
629+
630+ This function has to deal with dist-info based distributions and egg based distributions
631+ (which are still very much in use for "editable" installs).
632+
633+ Here are the file names as seen in a dist-info based distribution:
634+
635+ pytest_mock/__init__.py
636+ pytest_mock/_version.py
637+ pytest_mock/plugin.py
638+ pytest_mock.egg-info/PKG-INFO
639+
640+ Here are the file names as seen in an egg based distribution:
641+
642+ src/pytest_mock/__init__.py
643+ src/pytest_mock/_version.py
644+ src/pytest_mock/plugin.py
645+ src/pytest_mock.egg-info/PKG-INFO
646+ LICENSE
647+ setup.py
648+
649+ We have to take in account those two distribution flavors in order to determine which
650+ names should be considered for assertion rewriting.
651+
652+ More information:
653+ https://github.com/pytest-dev/pytest-mock/issues/167
654+ """
655+ package_files = list (package_files )
656+ seen_some = False
625657 for fn in package_files :
626658 is_simple_module = "/" not in fn and fn .endswith (".py" )
627659 is_package = fn .count ("/" ) == 1 and fn .endswith ("__init__.py" )
628660 if is_simple_module :
629661 module_name , _ = os .path .splitext (fn )
630- yield module_name
662+ # we ignore "setup.py" at the root of the distribution
663+ if module_name != "setup" :
664+ seen_some = True
665+ yield module_name
631666 elif is_package :
632667 package_name = os .path .dirname (fn )
668+ seen_some = True
633669 yield package_name
634670
671+ if not seen_some :
672+ # at this point we did not find any packages or modules suitable for assertion
673+ # rewriting, so we try again by stripping the first path component (to account for
674+ # "src" based source trees for example)
675+ # this approach lets us have the common case continue to be fast, as egg-distributions
676+ # are rarer
677+ new_package_files = []
678+ for fn in package_files :
679+ parts = fn .split ("/" )
680+ new_fn = "/" .join (parts [1 :])
681+ if new_fn :
682+ new_package_files .append (new_fn )
683+ if new_package_files :
684+ for _module in _iter_rewritable_modules (new_package_files ):
685+ yield _module
686+
635687
636688class Config (object ):
637689 """ access to configuration values, pluginmanager and plugin hooks. """
0 commit comments