devenv: Set PYTHONPATH where we install python modules

1 files changed, 44 insertions, 6 deletions

diff --git a/mesonbuild/modules/python.py b/mesonbuild/modules/python.py
index 5b91efc..f149ca1 100644
--- a/mesonbuild/modules/python.py
+++ b/mesonbuild/modules/python.py
@@ -23,7 +23,7 @@ from . import ExtensionModule
 from .. import mesonlib
 from .. import mlog
 from ..coredata import UserFeatureOption
-from ..build import known_shmod_kwargs
+from ..build import known_shmod_kwargs, EnvironmentVariables
 from ..dependencies import DependencyMethods, PkgConfigDependency, NotFoundDependency, SystemDependency, ExtraFrameworkDependency
 from ..dependencies.base import process_method_kw
 from ..environment import detect_cpu_family
@@ -45,6 +45,7 @@ if T.TYPE_CHECKING:
     from ..environment import Environment
     from ..interpreter import Interpreter
     from ..interpreterbase.interpreterbase import TYPE_var, TYPE_kwargs
+    from ..backends import InstallData
 
     from typing_extensions import TypedDict
 
@@ -389,6 +390,21 @@ class PythonExternalProgram(ExternalProgram):
             'variables': {},
             'version': '0.0',
         }
+        self.devenv_pythonpath: T.Set[str] = set()
+
+    def add_devenv_pythonpath(self, basedir: str, subdir: str, install_subdir: str) -> None:
+        # If we install python module into 'foo/bar' subdir, we need the last 2
+        # parts of source dir to be ['foo', 'bar'] and set PYTHONPATH
+        # pointing grandparent directory. That way scripts will be able to
+        # `import foo.bar.something` just like when the are installed.
+        # If the source tree layout does not match installed layout there is
+        # nothing we can do.
+        install_subdir_parts = Path(install_subdir).parts
+        subdir_parts = Path(subdir).parts
+        if subdir_parts[-len(install_subdir_parts):] == install_subdir_parts:
+            pypath = os.path.join(basedir, *subdir_parts[:-len(install_subdir_parts)])
+            self.devenv_pythonpath.add(pypath)
+            print('done', pypath)
 
     def _check_version(self, version: str) -> bool:
         if self.name == 'python2':
@@ -505,8 +521,10 @@ class PythonInstallation(ExternalProgramHolder):
             subdir = kwargs.pop('subdir', '')
             if not isinstance(subdir, str):
                 raise InvalidArguments('"subdir" argument must be a string.')
-
             kwargs['install_dir'] = os.path.join(self.platlib_install_path, subdir)
+            self.held_object.add_devenv_pythonpath(
+                self.interpreter.environment.get_build_dir(),
+                self.interpreter.subdir, subdir)
 
         # On macOS and some Linux distros (Debian) distutils doesn't link
         # extensions against libpython. We call into distutils and mirror its
@@ -561,11 +579,19 @@ class PythonInstallation(ExternalProgramHolder):
     def install_sources_method(self, args: T.Tuple[T.List[T.Union[str, mesonlib.File]]],
                                kwargs: 'PyInstallKw') -> 'Data':
         tag = kwargs['install_tag'] or 'runtime'
-        return self.interpreter.install_data_impl(
-            self.interpreter.source_strings_to_files(args[0]),
-            self._get_install_dir_impl(kwargs['pure'], kwargs['subdir']),
+        pure = kwargs['pure']
+        sources = self.interpreter.source_strings_to_files(args[0])
+        install_subdir = kwargs['subdir']
+        install_dir = self._get_install_dir_impl(pure, install_subdir)
+        builddir = self.interpreter.environment.get_build_dir()
+        srcdir = self.interpreter.environment.get_source_dir()
+        for src in sources:
+            basedir = builddir if src.is_built else srcdir
+            subdir = os.path.dirname(src.relative_name())
+            self.held_object.add_devenv_pythonpath(basedir, subdir, install_subdir)
+        return self.interpreter.install_data_impl(sources, install_dir,
             mesonlib.FileMode(), rename=None, tag=tag, install_data_type='python',
-            install_dir_name=self._get_install_dir_name_impl(kwargs['pure'], kwargs['subdir']))
+            install_dir_name=self._get_install_dir_name_impl(pure, install_subdir))
 
     @noPosargs
     @typed_kwargs('python_installation.install_dir', _PURE_KW, _SUBDIR_KW)
@@ -642,6 +668,18 @@ class PythonModule(ExtensionModule):
             'find_installation': self.find_installation,
         })
 
+    def get_devenv(self) -> T.Optional[EnvironmentVariables]:
+        pythonpath = set()
+        for python in self.installations.values():
+            version = python.info['version']
+            if mesonlib.version_compare(version, '>=3.0'):
+                pythonpath |= python.devenv_pythonpath
+        if pythonpath:
+            env = EnvironmentVariables()
+            env.prepend('PYTHONPATH', list(pythonpath))
+            return env
+        return None
+
     # https://www.python.org/dev/peps/pep-0397/
     @staticmethod
     def _get_win_pythonpath(name_or_path: str) -> T.Optional[str]: