Repository: pex-tool/pex

Issue:
Add option to include a Python package/module in the Pex
The `-D DIR` option adds directory entries in `DIR` at the root of the Pex import namespace.

This is awkward given a typical workspace layout like

```
.github/workflows/...
pyproject.toml
README.md
.venv/...
myapp/__init__.py
```

as giving `-D .` includes lots of junk that is also in the workspace (up to and including whole virtualenvs) and giving `-D myapp` does not create a `myapp` package in the Pex, it puts `__init__.py` into the root instead.

A workspace layout with a `src/` directory works but it means that a plain `python` command cannot import `myapp`.

I would propose adding an option (say `-P`) to include a single Python package or module in the PEX.

With this option:

* `-P myapp` would copy the `myapp` package into the PEX and make it importable as `myapp`
* `-P myapp.py` would copy the `myapp.py` module into the PEX and make it importable as `myapp`


Discussion:
So, Pex itself has this problem, but it, like most Python projects, has a project metadata file; so Pex uses `pex . -c pex -o pex.pex` to build itself as a PEX. Even though the `.` dir contains the `pex` package it also contains the `scripts` dir, etc but the built `pex.pex` contains none of this extraneous stuff. In other words, the `-D` was added late in Pex history to accommodate uses (Pants was the motivating case) that did not necessarily have a traditional project metadata file (`setup.py` or `pyproject.toml`).

Does the existing project metadata support suffice?
Fwiw the actual production release build of the Pex PEX released to GitHub releases uses this:
https://github.com/pantsbuild/pex/blob/382d7aa1fdb89072cb33863565ff67e45000a4f3/scripts/package.py#L28-L59
Ah, I'm seeking to avoid needing our applications to build with setuptools/flit/poetry. This is in a corporate environment, for context.

We've tended to distinguish between libraries and applications. Libraries publish wheels and sdists so they need a package build tool - but the bar is generally going to be higher, ensuring they work across Python versions, dependency versions and so on.

But applications just run with a `requirements.txt` and some Python code. That feels like less for user to think about, particularly when it comes to versioning. For deployment we just have to provide a virtualenv with the requirements. We're exploring Pex for this CLI applications. For services we will just build Docker images.

I now have a workaround that I'm happy with, which is to export the package tree into a temp dir using git. Unlike a plain copy, which only exists to overcome pex's command line limitations, it has the desirable side-effect of excluding untracked files and uncommitted changes, so that seems to justify it a little more:

```bash
TMPDIR="$(mktemp -d)"
trap 'rm -rf -- "$TMPDIR"' EXIT
git archive HEAD myapp | tar x -C "$TMPDIR"
pex -D "$TMPDIR"
```
> Ah, I'm seeking to avoid needing our applications to build with setuptools/flit/poetry. This is in a corporate environment, for context.

Aha, Ok. Your OP included a `pyproject.toml` mention and I assumed. Well, you are carving a niche I have not heard about use of yet which is large corp / project that uses Pex but doesn't use Pants. Pants would allow you this control when generating a PEX.

So your use case is not unreasonable (Pants rolls this way - not forcing and in fact even discouraging a python distribution for anything but to-be-published libraries). The trick is just nailing down the interface here. I could imagine an extension of the current syntax like `-D myapp=myapp`, where the `=` is the trigger to take the LHS as the directory to use and the RHS as the prefix to give files in that directory.
We have a similar use for packaging pyspark jobs with pip dependencies. Packaging the pip dependencies is fine, but including the job module (and other modules it imports) is cumbersome. We don't want to create package metadata for each job, we just want to include a few directories in the resulting pex file and have them being importable at runtime.

The `-D myapp=myapp` syntax would be fantastic.
@khaledh thanks for chiming in with a second use case. @lordmauve I've gone with your suggestion roughly. `-P foo` to grab package `foo` and all sub-packages in the CWD with support for `-P foo.bar` as well (drill down) and, finally, `-P foo@src` for functionality ~redundant with `-D src` when foo is the only top-level package, but otherwise providing more selection flexibility when it is not. Likewise there is a `-M` option to pick out individual modules.

Resolving patch:
diff --git a/pex/bin/pex.py b/pex/bin/pex.py
--- a/pex/bin/pex.py
+++ b/pex/bin/pex.py
@@ -23,7 +23,7 @@
     global_environment,
     register_global_arguments,
 )
-from pex.common import die, safe_mkdtemp
+from pex.common import die, filter_pyc_dirs, filter_pyc_files, safe_mkdtemp
 from pex.enum import Enum
 from pex.inherit_path import InheritPath
 from pex.interpreter_constraints import InterpreterConstraints
@@ -48,9 +48,14 @@
 
 if TYPE_CHECKING:
     from argparse import Namespace
-    from typing import Dict, List, Optional
+    from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple
+
+    import attr  # vendor:skip
 
     from pex.resolve.resolver_options import ResolverConfiguration
+else:
+    from pex.third_party import attr
+
 
 CANNOT_SETUP_INTERPRETER = 102
 INVALID_OPTIONS = 103
@@ -461,6 +466,149 @@ def __call__(self, parser, namespace, value, option_str=None):
         setattr(namespace, self.dest, seed)
 
 
+@attr.s(frozen=True)
+class PythonSource(object):
+    @classmethod
+    def parse(cls, name):
+        # type: (str) -> PythonSource
+        subdir = None
+        parts = name.split("@", 1)
+        if len(parts) == 2:
+            name, subdir = parts
+        return cls(name=name, subdir=subdir)
+
+    name = attr.ib()  # type: str
+    subdir = attr.ib(default=None)  # type: Optional[str]
+
+    def iter_files(self):
+        # type: () -> Iterator[Tuple[str, str]]
+        components = self.name.split(".")
+        parent_package_dirs = components[:-1]
+        source = components[-1]
+
+        package_path = [self.subdir] if self.subdir else []  # type: List[str]
+        for package_dir in parent_package_dirs:
+            package_path.append(package_dir)
+            package_file_src = os.path.join(*(package_path + ["__init__.py"]))
+            if os.path.exists(package_file_src):
+                package_file_dst = (
+                    os.path.relpath(package_file_src, self.subdir)
+                    if self.subdir
+                    else package_file_src
+                )
+                yield package_file_src, package_file_dst
+
+        for src, dst in self._iter_source_files(package_path, source):
+            yield src, dst
+
+    def _iter_source_files(
+        self,
+        parent_package_path,  # type: List[str]
+        source,  # type: str
+    ):
+        # type: (...) -> Iterator[Tuple[str, str]]
+        raise NotImplementedError()
+
+
+class Package(PythonSource):
+    def _iter_source_files(
+        self,
+        parent_package_path,  # type: List[str]
+        source,  # type: str
+    ):
+        # type: (...) -> Iterator[Tuple[str, str]]
+        package_dir = os.path.join(*(parent_package_path + [source]))
+        for root, dirs, files in os.walk(package_dir):
+            dirs[:] = list(filter_pyc_dirs(dirs))
+            for f in filter_pyc_files(files):
+                src = os.path.join(root, f)
+                dst = os.path.relpath(src, self.subdir) if self.subdir else src
+                yield src, dst
+
+
+class Module(PythonSource):
+    def _iter_source_files(
+        self,
+        parent_package_path,  # type: List[str]
+        source,  # type: str
+    ):
+        # type: (...) -> Iterator[Tuple[str, str]]
+        module_src = os.path.join(*(parent_package_path + ["{module}.py".format(module=source)]))
+        module_dest = os.path.relpath(module_src, self.subdir) if self.subdir else module_src
+        yield module_src, module_dest
+
+
+def configure_clp_sources(parser):
+    # type: (ArgumentParser) -> None
+
+    parser.add_argument(
+        "-D",
+        "--sources-directory",
+        dest="sources_directory",
+        metavar="DIR",
+        default=[],
+        type=str,
+        action="append",
+        help=(
+            "Add a directory containing sources and/or resources to be packaged into the generated "
+            ".pex file. This option can be used multiple times."
+        ),
+    )
+
+    parser.add_argument(
+        "-R",
+        "--resources-directory",
+        dest="resources_directory",
+        metavar="DIR",
+        default=[],
+        type=str,
+        action="append",
+        help=(
+            "Add resources directory to be packaged into the generated .pex file."
+            " This option can be used multiple times. DEPRECATED: Use -D/--sources-directory "
+            "instead."
+        ),
+    )
+
+    parser.add_argument(
+        "-P",
+        "--package",
+        dest="packages",
+        metavar="PACKAGE_SPEC",
+        default=[],
+        type=Package.parse,
+        action="append",
+        help=(
+            "Add a package and all its sub-packages to the generated .pex file. The package is "
+            "expected to be found relative to the the current directory. If the package is housed "
+            "in a subdirectory, indicate that by appending `@<subdirectory>`. For example, to add "
+            "the top-level package `foo` housed in the current directory, use `-P foo`. If the "
+            "top-level `foo` package is in the `src` subdirectory use `-P foo@src`. If you wish to "
+            "just use the `foo.bar` package in the `src` subdirectory, use `-P foo.bar@src`. This "
+            "option can be used multiple times."
+        ),
+    )
+
+    parser.add_argument(
+        "-M",
+        "--module",
+        dest="modules",
+        metavar="MODULE_SPEC",
+        default=[],
+        type=Module.parse,
+        action="append",
+        help=(
+            "Add an individual module to the generated .pex file. The module is expected to be "
+            "found relative to the the current directory. If the module is housed in a "
+            "subdirectory, indicate that by appending `@<subdirectory>`. For example, to add the "
+            "top-level module `foo` housed in the current directory, use `-M foo`. If the "
+            "top-level `foo` module is in the `src` subdirectory use `-M foo@src`. If you wish to "
+            "just use the `foo.bar` module in the `src` subdirectory, use `-M foo.bar@src`. This "
+            "option can be used multiple times."
+        ),
+    )
+
+
 def configure_clp():
     # type: () -> ArgumentParser
     usage = (
@@ -504,35 +652,7 @@ def configure_clp():
         help="The name of a file to be included as the preamble for the generated .pex file",
     )
 
-    parser.add_argument(
-        "-D",
-        "--sources-directory",
-        dest="sources_directory",
-        metavar="DIR",
-        default=[],
-        type=str,
-        action="append",
-        help=(
-            "Add a directory containing sources and/or resources to be packaged into the generated "
-            ".pex file. This option can be used multiple times."
-        ),
-    )
-
-    parser.add_argument(
-        "-R",
-        "--resources-directory",
-        dest="resources_directory",
-        metavar="DIR",
-        default=[],
-        type=str,
-        action="append",
-        help=(
-            "Add resources directory to be packaged into the generated .pex file."
-            " This option can be used multiple times. DEPRECATED: Use -D/--sources-directory "
-            "instead."
-        ),
-    )
-
+    configure_clp_sources(parser)
     requirement_options.register(parser)
 
     parser.add_argument(
@@ -580,6 +700,24 @@ def configure_clp():
     return parser
 
 
+def _iter_directory_sources(directories):
+    # type: (Iterable[str]) -> Iterator[Tuple[str, str]]
+    for directory in directories:
+        src_dir = os.path.normpath(directory)
+        for root, _, files in os.walk(src_dir):
+            for f in files:
+                src_file_path = os.path.join(root, f)
+                dst_path = os.path.relpath(src_file_path, src_dir)
+                yield src_file_path, dst_path
+
+
+def _iter_python_sources(python_sources):
+    # type: (Iterable[PythonSource]) -> Iterator[Tuple[str, str]]
+    for python_source in python_sources:
+        for src, dst in python_source.iter_files():
+            yield src, dst
+
+
 def build_pex(
     requirement_configuration,  # type: RequirementConfiguration
     resolver_configuration,  # type: ResolverConfiguration
@@ -626,16 +764,16 @@ def build_pex(
             "dependency cache."
         )
 
-    directories = OrderedSet(
-        options.sources_directory + options.resources_directory
-    )  # type: OrderedSet[str]
-    for directory in directories:
-        src_dir = os.path.normpath(directory)
-        for root, _, files in os.walk(src_dir):
-            for f in files:
-                src_file_path = os.path.join(root, f)
-                dst_path = os.path.relpath(src_file_path, src_dir)
-                pex_builder.add_source(src_file_path, dst_path)
+    seen = set()  # type: Set[Tuple[str, str]]
+    for src, dst in itertools.chain(
+        _iter_directory_sources(
+            OrderedSet(options.sources_directory + options.resources_directory)
+        ),
+        _iter_python_sources(OrderedSet(options.packages + options.modules)),
+    ):
+        if (src, dst) not in seen:
+            pex_builder.add_source(src, dst)
+            seen.add((src, dst))
 
     pex_info = pex_builder.info
     pex_info.inject_env = dict(options.inject_env)


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -394,6 +394,7 @@ def run_pex_command(
     env=None,  # type: Optional[Dict[str, str]]
     python=None,  # type: Optional[str]
     quiet=False,  # type: bool
+    cwd=None,  # type: Optional[str]
 ):
     # type: (...) -> IntegResults
     """Simulate running pex command for integration testing.
@@ -404,7 +405,7 @@ def run_pex_command(
     """
     cmd = create_pex_command(args, python=python, quiet=quiet)
     process = Executor.open_process(
-        cmd=cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        cmd=cmd, env=env, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
     )
     output, error = process.communicate()
     return IntegResults(output.decode("utf-8"), error.decode("utf-8"), process.returncode)
diff --git a/tests/integration/test_issue_2134.py b/tests/integration/test_issue_2134.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_issue_2134.py
@@ -0,0 +1,212 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os.path
+import subprocess
+import sys
+from typing import Iterable, Optional
+
+import pytest
+
+from pex import layout
+from pex.common import open_zip, touch
+from pex.testing import run_pex_command
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    import attr  # vendor:skip
+else:
+    from pex.third_party import attr
+
+
+@attr.s(frozen=True)
+class SourceTree(object):
+    base_dir = attr.ib()  # type: str
+
+    def create(self, offset=None):
+        # type: (Optional[str]) -> str
+
+        project_dir = os.path.join(self.base_dir, "project")
+        source_root = os.path.join(project_dir, offset) if offset else project_dir
+
+        touch(os.path.join(source_root, "top_level_module.py"))
+
+        touch(os.path.join(source_root, "top_level_package_classic", "__init__.py"))
+        touch(os.path.join(source_root, "top_level_package_classic", "module.py"))
+        touch(os.path.join(source_root, "top_level_package_classic", "sub", "__init__.py"))
+        touch(os.path.join(source_root, "top_level_package_classic", "sub", "foo.py"))
+        touch(os.path.join(source_root, "top_level_package_classic", "sub", "bar.py"))
+
+        # N.B.: This test will run against a range of interpreters, some supporting PEP420 and some
+        # not, so we never import code in the tests that use this fixture and instead just check
+        # file lists,
+        touch(os.path.join(source_root, "top_level_package_pep420", "module.py"))
+        touch(os.path.join(source_root, "top_level_package_pep420", "sub", "module.py"))
+
+        touch(os.path.join(source_root, "top_level_package_mixed", "module.py"))
+        touch(os.path.join(source_root, "top_level_package_mixed", "sub", "__init__.py"))
+        touch(os.path.join(source_root, "top_level_package_mixed", "sub", "module.py"))
+
+        subprocess.check_call(args=[sys.executable, "-m", "compileall", source_root])
+
+        return project_dir
+
+    def assert_sources(
+        self,
+        pex_args,  # type: Iterable[str]
+        expected_sources,  # type: Iterable[str]
+        offset=None,  # type: Optional[str]
+    ):
+        # type: (...) -> None
+        pex = os.path.join(self.base_dir, "pex")
+        args = ["-o", pex]
+        args.extend(pex_args)
+        run_pex_command(args=args, cwd=self.create(offset=offset)).assert_success()
+
+        with open_zip(pex) as zf:
+            actual_sources = [
+                f
+                for f in zf.namelist()
+                if not f.startswith((layout.BOOTSTRAP_DIR, layout.DEPS_DIR, "__pex__"))
+                and f not in (layout.PEX_INFO_PATH, "__main__.py")
+            ]
+
+        assert sorted(expected_sources) == sorted(actual_sources)
+
+
+@pytest.fixture
+def source_tree(tmpdir):
+    # type: (Any) -> SourceTree
+    return SourceTree(str(tmpdir))
+
+
+def test_add_top_level_package(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-P", "top_level_package_classic"],
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/module.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/foo.py",
+            "top_level_package_classic/sub/bar.py",
+        ],
+    )
+
+
+def test_add_sub_package(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-P", "top_level_package_classic.sub"],
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/foo.py",
+            "top_level_package_classic/sub/bar.py",
+        ],
+    )
+
+
+def test_add_sub_package_pep_420(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-P", "top_level_package_pep420.sub"],
+        expected_sources=[
+            "top_level_package_pep420/",
+            "top_level_package_pep420/sub/",
+            "top_level_package_pep420/sub/module.py",
+        ],
+    )
+
+
+def test_add_sub_package_mixed(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-P", "top_level_package_mixed.sub"],
+        expected_sources=[
+            "top_level_package_mixed/",
+            "top_level_package_mixed/sub/",
+            "top_level_package_mixed/sub/__init__.py",
+            "top_level_package_mixed/sub/module.py",
+        ],
+    )
+
+
+def test_add_package_offset(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-P", "top_level_package_classic.sub@src"],
+        offset="src",
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/foo.py",
+            "top_level_package_classic/sub/bar.py",
+        ],
+    )
+
+
+def test_add_top_level_module(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-M", "top_level_module"], expected_sources=["top_level_module.py"]
+    )
+
+
+def test_add_module_in_package(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=["-M", "top_level_package_classic.sub.foo"],
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/foo.py",
+        ],
+    )
+
+
+def test_add_module_offset(source_tree):
+    # type: (SourceTree) -> None
+    offset = os.path.join("src", "python")
+    source_tree.assert_sources(
+        pex_args=["-M", "top_level_package_classic.sub.bar@{offset}".format(offset=offset)],
+        offset=offset,
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/bar.py",
+        ],
+    )
+
+
+def test_overlap(source_tree):
+    # type: (SourceTree) -> None
+    source_tree.assert_sources(
+        pex_args=[
+            "-M",
+            "top_level_package_classic.sub.foo@src",
+            "-P",
+            "top_level_package_classic.sub@src",
+        ],
+        offset="src",
+        expected_sources=[
+            "top_level_package_classic/",
+            "top_level_package_classic/__init__.py",
+            "top_level_package_classic/sub/",
+            "top_level_package_classic/sub/__init__.py",
+            "top_level_package_classic/sub/foo.py",
+            "top_level_package_classic/sub/bar.py",
+        ],
+    )