Add helper to build in-tree extensions. (#2831)

For single-file extensions, a convenient pattern offered by cython
is to place the source files directly in the python source tree
(`foo/__init__.py`, `foo/ext.pyx`), deriving the package names from
their filesystem location.  Adapt this pattern for pybind11, using an
`intree_extensions` helper, which should be thought of as the moral
equivalent to `cythonize`.

Differences with cythonize: I chose not to include globbing support
(`intree_extensions(glob.glob("**/*.cpp"))` seems sufficient), nor to
provide extension-customization kwargs (directly setting the attributes
on the resulting Pybind11Extension objects seems sufficient).

We could choose to have `intree_extension` (singular instead) and make
users write `[*map(intree_extension, glob.glob("**/*.cpp"))]`; no strong
opinion here.

Co-authored-by: Aaron Gokaslan <skylion.aaron@gmail.com>

docs/compiling.rst

@@ -70,6 +70,19 @@ that is supported via a ``build_ext`` command override; it will only affect
         ext_modules=ext_modules
     )
 
+If you have single-file extension modules that are directly stored in the
+Python source tree (``foo.cpp`` in the same directory as where a ``foo.py``
+would be located), you can also generate ``Pybind11Extensions`` using
+``setup_helpers.intree_extensions``: ``intree_extensions(["path/to/foo.cpp",
+...])`` returns a list of ``Pybind11Extensions`` which can be passed to
+``ext_modules``, possibly after further customizing their attributes
+(``libraries``, ``include_dirs``, etc.).  By doing so, a ``foo.*.so`` extension
+module will be generated and made available upon installation.
+
+``intree_extension`` will automatically detect if you are using a ``src``-style
+layout (as long as no namespace packages are involved), but you can also
+explicitly pass ``package_dir`` to it (as in ``setuptools.setup``).
+
 Since pybind11 does not require NumPy when building, a light-weight replacement
 for NumPy's parallel compilation distutils tool is included. Use it like this:
 

pybind11/setup_helpers.py

@@ -303,6 +303,42 @@ class build_ext(_build_ext):  # noqa: N801
         _build_ext.build_extensions(self)
 
 
+def intree_extensions(paths, package_dir=None):
+    """
+    Generate Pybind11Extensions from source files directly located in a Python
+    source tree.
+
+    ``package_dir`` behaves as in ``setuptools.setup``.  If unset, the Python
+    package root parent is determined as the first parent directory that does
+    not contain an ``__init__.py`` file.
+    """
+    exts = []
+    for path in paths:
+        if package_dir is None:
+            parent, _ = os.path.split(path)
+            while os.path.exists(os.path.join(parent, "__init__.py")):
+                parent, _ = os.path.split(parent)
+            relname, _ = os.path.splitext(os.path.relpath(path, parent))
+            qualified_name = relname.replace(os.path.sep, ".")
+            exts.append(Pybind11Extension(qualified_name, [path]))
+        else:
+            found = False
+            for prefix, parent in package_dir.items():
+                if path.startswith(parent):
+                    found = True
+                    relname, _ = os.path.splitext(os.path.relpath(path, parent))
+                    qualified_name = relname.replace(os.path.sep, ".")
+                    if prefix:
+                        qualified_name = prefix + "." + qualified_name
+                    exts.append(Pybind11Extension(qualified_name, [path]))
+            if not found:
+                raise ValueError(
+                    "path {} is not a child of any of the directories listed "
+                    "in 'package_dir' ({})".format(path, package_dir)
+                )
+    return exts
+
+
 def naive_recompile(obj, src):
     """
     This will recompile only if the source file changes. It does not check

pybind11/setup_helpers.pyi

@@ -1,7 +1,7 @@
 # IMPORTANT: Should stay in sync with setup_helpers.py (mostly checked by CI /
 # pre-commit).
 
-from typing import Any, Callable, Iterator, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Type, TypeVar, Union
 from types import TracebackType
 
 from distutils.command.build_ext import build_ext as _build_ext  # type: ignore
@@ -33,6 +33,9 @@ def auto_cpp_level(compiler: distutils.ccompiler.CCompiler) -> Union[int, str]:
 class build_ext(_build_ext):  # type: ignore
     def build_extensions(self) -> None: ...
 
+def intree_extensions(
+    paths: Iterator[str], package_dir: Optional[Dict[str, str]] = None
+) -> List[Pybind11Extension]: ...
 def no_recompile(obj: str, src: str) -> bool: ...
 def naive_recompile(obj: str, src: str) -> bool: ...
 

tests/extra_setuptools/test_setuphelper.py

@@ -99,3 +99,45 @@ def test_simple_setup_py(monkeypatch, tmpdir, parallel, std):
     subprocess.check_call(
         [sys.executable, "test.py"], stdout=sys.stdout, stderr=sys.stderr
     )
+
+
+def test_intree_extensions(monkeypatch, tmpdir):
+    monkeypatch.syspath_prepend(MAIN_DIR)
+
+    from pybind11.setup_helpers import intree_extensions
+
+    monkeypatch.chdir(tmpdir)
+    root = tmpdir
+    root.ensure_dir()
+    subdir = root / "dir"
+    subdir.ensure_dir()
+    src = subdir / "ext.cpp"
+    src.ensure()
+    (ext,) = intree_extensions([src.relto(tmpdir)])
+    assert ext.name == "ext"
+    subdir.ensure("__init__.py")
+    (ext,) = intree_extensions([src.relto(tmpdir)])
+    assert ext.name == "dir.ext"
+
+
+def test_intree_extensions_package_dir(monkeypatch, tmpdir):
+    monkeypatch.syspath_prepend(MAIN_DIR)
+
+    from pybind11.setup_helpers import intree_extensions
+
+    monkeypatch.chdir(tmpdir)
+    root = tmpdir / "src"
+    root.ensure_dir()
+    subdir = root / "dir"
+    subdir.ensure_dir()
+    src = subdir / "ext.cpp"
+    src.ensure()
+    (ext,) = intree_extensions([src.relto(tmpdir)], package_dir={"": "src"})
+    assert ext.name == "dir.ext"
+    (ext,) = intree_extensions([src.relto(tmpdir)], package_dir={"foo": "src"})
+    assert ext.name == "foo.dir.ext"
+    subdir.ensure("__init__.py")
+    (ext,) = intree_extensions([src.relto(tmpdir)], package_dir={"": "src"})
+    assert ext.name == "dir.ext"
+    (ext,) = intree_extensions([src.relto(tmpdir)], package_dir={"foo": "src"})
+    assert ext.name == "foo.dir.ext"