From 9d81ad40e71211a4d307395a14b94dd04cb9af8f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miro=20Hron=C4=8Dok?= <miro@hroncok.cz>
Date: Mon, 1 Nov 2021 15:53:43 +0100
Subject: [PATCH] %py(3)_check_import: Process .pth files in site(arch|lib)

Fixes https://bugzilla.redhat.com/show_bug.cgi?id=2018551
---
 import_all_modules.py            | 15 +++++++++++++++
 macros.python                    |  1 +
 macros.python3                   |  1 +
 python-rpm-macros.spec           |  2 ++
 tests/test_evals.py              |  1 +
 tests/test_import_all_modules.py | 33 ++++++++++++++++++++++++++++++++
 6 files changed, 53 insertions(+)

diff --git a/import_all_modules.py b/import_all_modules.py
index 5788b8c..3930236 100644
--- a/import_all_modules.py
+++ b/import_all_modules.py
@@ -3,7 +3,9 @@
 import argparse
 import importlib
 import fnmatch
+import os
 import re
+import site
 import sys
 
 from contextlib import contextmanager
@@ -139,6 +141,18 @@ def remove_unwanteds_from_sys_path():
         sys.path = old_sys_path
 
 
+def addsitedirs_from_environ():
+    '''Load directories from the _PYTHONSITE environment variable (separated by :)
+    and load the ones already present in sys.path via site.addsitedir()
+    to handle .pth files in them.
+
+    This is needed to properly import old-style namespace packages with nspkg.pth files.
+    See https://bugzilla.redhat.com/2018551 for a more detailed rationale.'''
+    for path in os.getenv('_PYTHONSITE', '').split(':'):
+        if path in sys.path:
+            site.addsitedir(path)
+
+
 def main(argv=None):
 
     cli_args = argparser().parse_args(argv)
@@ -149,6 +163,7 @@ def main(argv=None):
     modules = read_modules_from_all_args(cli_args)
 
     with remove_unwanteds_from_sys_path():
+        addsitedirs_from_environ()
         import_modules(modules)
 
 
diff --git a/macros.python b/macros.python
index 30c6fc3..f0d431d 100644
--- a/macros.python
+++ b/macros.python
@@ -79,6 +79,7 @@
 %py_check_import(e:tf:) %{expand:\\\
   PATH="%{buildroot}%{_bindir}:$PATH"\\\
   PYTHONPATH="${PYTHONPATH:-%{buildroot}%{python_sitearch}:%{buildroot}%{python_sitelib}}"\\\
+  _PYTHONSITE="%{buildroot}%{python_sitearch}:%{buildroot}%{python_sitelib}"\\\
   PYTHONDONTWRITEBYTECODE=1\\\
   %{lua:
   local command = "%{__python} "
diff --git a/macros.python3 b/macros.python3
index c004acf..0ccafd9 100644
--- a/macros.python3
+++ b/macros.python3
@@ -77,6 +77,7 @@
 %py3_check_import(e:tf:) %{expand:\\\
   PATH="%{buildroot}%{_bindir}:$PATH"\\\
   PYTHONPATH="${PYTHONPATH:-%{buildroot}%{python3_sitearch}:%{buildroot}%{python3_sitelib}}"\\\
+  _PYTHONSITE="%{buildroot}%{python3_sitearch}:%{buildroot}%{python3_sitelib}"\\\
   PYTHONDONTWRITEBYTECODE=1\\\
   %{lua:
   local command = "%{__python3} "
diff --git a/python-rpm-macros.spec b/python-rpm-macros.spec
index f990122..aaeddd9 100644
--- a/python-rpm-macros.spec
+++ b/python-rpm-macros.spec
@@ -150,6 +150,8 @@ install -m 755 brp-* %{buildroot}%{_rpmconfigdir}/redhat/
 Resolves: rhbz#2018809
 - Fix %%py_shebang_flags handling within %%py_check_import
 Resolves: rhbz#2018615
+- Process .pth files in buildroot's sitedirs in %%py_check_import
+Resolves: rhbz#2018551
 
 * Mon Oct 25 2021 Karolina Surma <ksurma@redhat.com> - 3.10-12
 - Introduce -f (read from file) option to %%py{3}_check_import
diff --git a/tests/test_evals.py b/tests/test_evals.py
index 7cec492..6a771ed 100644
--- a/tests/test_evals.py
+++ b/tests/test_evals.py
@@ -714,6 +714,7 @@ def test_py3_check_import(args, expected_args, __python3, lib):
     expected = textwrap.dedent(fr"""
         PATH="BUILDROOT/usr/bin:$PATH"
         PYTHONPATH="${{PYTHONPATH:-BUILDROOT/usr/{lib}/python{x_y}/site-packages:BUILDROOT/usr/lib/python{x_y}/site-packages}}"
+        _PYTHONSITE="BUILDROOT/usr/{lib}/python{x_y}/site-packages:BUILDROOT/usr/lib/python{x_y}/site-packages"
         PYTHONDONTWRITEBYTECODE=1
         {__python3 or '/usr/bin/python3'} -s RPMCONFIGDIR/redhat/import_all_modules.py {expected_args}
         """)
diff --git a/tests/test_import_all_modules.py b/tests/test_import_all_modules.py
index d60f6e4..52e1d7e 100644
--- a/tests/test_import_all_modules.py
+++ b/tests/test_import_all_modules.py
@@ -391,3 +391,36 @@ def test_module_list_from_relative_path(tmp_path, monkeypatch):
     ])
 
     assert 'wave' in sys.modules
+
+
+@pytest.mark.parametrize('arch_in_path', [True, False])
+def test_pth_files_are_read_from__PYTHONSITE(arch_in_path, tmp_path, monkeypatch, capsys):
+    sitearch = tmp_path / 'lib64'
+    sitearch.mkdir()
+    sitelib = tmp_path / 'lib'
+    sitelib.mkdir()
+
+    for where, word in (sitearch, "ARCH"), (sitelib, "LIB"), (sitelib, "MOD"):
+        module = where / f'print{word}.py'
+        module.write_text(f'print("{word}")')
+
+    pth_sitearch = sitearch / 'ARCH.pth'
+    pth_sitearch.write_text('import printARCH\n')
+
+    pth_sitelib = sitelib / 'LIB.pth'
+    pth_sitelib.write_text('import printLIB\n')
+
+    if arch_in_path:
+        sys.path.append(str(sitearch))
+    sys.path.append(str(sitelib))
+
+    # we always add sitearch to _PYTHONSITE
+    # but when not in sys.path, it should not be processed for .pth files
+    monkeypatch.setenv('_PYTHONSITE', f'{sitearch}:{sitelib}')
+
+    modules_main(['printMOD'])
+    out, err = capsys.readouterr()
+    if arch_in_path:
+        assert out == 'ARCH\nLIB\nMOD\n'
+    else:
+        assert out == 'LIB\nMOD\n'