pw_build: Optionally generate setup.py & nest protos

- Support the generated_setup argument for pw_python_package. This
  mirrors the package to the out directory and generates a setup.py
  for it there.
- Allow pw_proto_library targets to add their protos to an existing
  Python package rather than generating a protos-only package.
- Only reinstall --editable packages when setup.py changes rather than
  when any file changes.

Requires: pigweed-internal:10840
Change-Id: I35ed555c2667e60d844468eb614ce0a6f76c3d32
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/36504
Commit-Queue: Wyatt Hepler <hepler@google.com>
Reviewed-by: Keir Mierle <keir@google.com>
Reviewed-by: Joe Ethier <jethier@google.com>
Reviewed-by: Alexei Frolov <frolv@google.com>
diff --git a/docs/python_build.rst b/docs/python_build.rst
index db194f7..4bc4fcd 100644
--- a/docs/python_build.rst
+++ b/docs/python_build.rst
@@ -302,11 +302,6 @@
 source tree is incomplete; the final Python package, including protobufs, is
 generated in the output directory.
 
-.. admonition:: Under construction
-
-  Protobuf modules are currently only generated as standalone packages. Support
-  for generating protobuf modules into existing packages will be added.
-
 Generating setup.py
 -------------------
 The ``pw_python_package`` target in the ``BUILD.gn`` duplicates much of the
diff --git a/pw_build/py/BUILD.gn b/pw_build/py/BUILD.gn
index 5287fb5..ea96dc4 100644
--- a/pw_build/py/BUILD.gn
+++ b/pw_build/py/BUILD.gn
@@ -23,6 +23,7 @@
     "pw_build/copy_from_cipd.py",
     "pw_build/error.py",
     "pw_build/exec.py",
+    "pw_build/generate_python_package.py",
     "pw_build/generate_python_package_gn.py",
     "pw_build/generated_tests.py",
     "pw_build/host_tool.py",
diff --git a/pw_build/py/pw_build/generate_python_package.py b/pw_build/py/pw_build/generate_python_package.py
new file mode 100644
index 0000000..199cd45
--- /dev/null
+++ b/pw_build/py/pw_build/generate_python_package.py
@@ -0,0 +1,202 @@
+# Copyright 2021 The Pigweed Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+"""Script that invokes protoc to generate code for .proto files."""
+
+import argparse
+from collections import defaultdict
+import json
+from pathlib import Path
+import sys
+import textwrap
+from typing import Dict, List, Set, TextIO
+
+try:
+    from pw_build.mirror_tree import mirror_paths
+except ImportError:
+    # Append this path to the module search path to allow running this module
+    # before the pw_build package is installed.
+    sys.path.append(str(Path(__file__).resolve().parent.parent))
+    from pw_build.mirror_tree import mirror_paths
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+
+    parser.add_argument('--file-list',
+                        required=True,
+                        type=argparse.FileType('r'),
+                        help='A list of files to copy')
+    parser.add_argument('--file-list-root',
+                        required=True,
+                        type=argparse.FileType('r'),
+                        help='A file with the root of the file list')
+    parser.add_argument('--label', help='Label for this Python package')
+    parser.add_argument('--proto-library',
+                        default='',
+                        help='Name of proto library nested in this package')
+    parser.add_argument('--proto-library-file',
+                        type=Path,
+                        help="File with the proto library's name")
+    parser.add_argument('--root',
+                        required=True,
+                        type=Path,
+                        help='The base directory for the Python package')
+    parser.add_argument('--setup-json',
+                        required=True,
+                        type=argparse.FileType('r'),
+                        help='setup.py keywords as JSON')
+    parser.add_argument('--module-as-package',
+                        action='store_true',
+                        help='Generate an __init__.py that imports everything')
+    parser.add_argument('files',
+                        type=Path,
+                        nargs='+',
+                        help='Relative paths to the files in the package')
+    return parser.parse_args()
+
+
+def _check_nested_protos(label: str, proto_library_file: Path,
+                         proto_library: str) -> None:
+    """Checks that the proto library refers to this package; returns error."""
+    error = 'not set'
+
+    if proto_library_file.exists():
+        proto_label = proto_library_file.read_text().strip()
+        if proto_label == label:
+            return
+
+        if proto_label:
+            error = f'set to {proto_label}'
+
+    raise ValueError(
+        f"{label}'s 'proto_library' is set to {proto_library}, but that "
+        f"target's 'python_package' is {error}. Set {proto_library}'s "
+        f"'python_package' to {label}.")
+
+
+def _collect_all_files(files: List[Path], root: Path, file_list: TextIO,
+                       file_list_root: TextIO) -> Dict[str, Set[str]]:
+    """Collects files in output dir, adds to files; returns package_data."""
+    root.mkdir(exist_ok=True)
+
+    other_files = [Path(p.rstrip()) for p in file_list]
+    other_files_root = Path(file_list_root.read().rstrip())
+
+    # Mirror the proto files to this package.
+    files += mirror_paths(other_files_root, other_files, root)
+
+    # Find all subpackages, including empty ones.
+    subpackages: Set[Path] = set()
+    for file in (f.relative_to(root) for f in files):
+        subpackages.update(root / path for path in file.parents)
+    subpackages.remove(root)
+
+    # Make sure there are __init__.py and py.typed files for each subpackage.
+    for pkg in subpackages:
+        for file in (pkg / name for name in ['__init__.py', 'py.typed']):
+            file.touch()
+            files.append(file)
+
+    pkg_data: Dict[str, Set[str]] = defaultdict(set)
+
+    # Add all non-source files to package data.
+    for file in (f for f in files if f.suffix != '.py'):
+        pkg = root / file.parent
+        package_name = pkg.relative_to(root).as_posix().replace('/', '.')
+        pkg_data[package_name].add(file.name)
+
+    return pkg_data
+
+
+_SETUP_PY_FILE = '''\
+# Generated file. Do not modify.
+# pylint: skip-file
+
+import setuptools  # type: ignore
+
+setuptools.setup(
+{keywords}
+)
+'''
+
+
+def _generate_setup_py(pkg_data: dict, setup_json: TextIO) -> str:
+    setup_keywords = dict(
+        packages=list(pkg_data),
+        package_data={pkg: list(files)
+                      for pkg, files in pkg_data.items()},
+    )
+
+    specified_keywords = json.load(setup_json)
+
+    assert not any(kw in specified_keywords for kw in setup_keywords), (
+        'Generated packages may not specify "packages" or "package_data"')
+    setup_keywords.update(specified_keywords)
+
+    return _SETUP_PY_FILE.format(keywords='\n'.join(
+        f'    {k}={v!r},' for k, v in setup_keywords.items()))
+
+
+def _import_module_in_package_init(all_files: List[Path]) -> None:
+    """Generates an __init__.py that imports the module.
+
+    This makes an individual module usable as a package. This is used for proto
+    modules.
+    """
+    sources = [
+        f for f in all_files if f.suffix == '.py' and f.name != '__init__.py'
+    ]
+    assert len(sources) == 1, (
+        'Module as package expects a single .py source file')
+
+    source, = sources
+    source.parent.joinpath('__init__.py').write_text(
+        f'from {source.stem}.{source.stem} import *\n')
+
+
+def main(files: List[Path],
+         root: Path,
+         file_list: TextIO,
+         file_list_root: TextIO,
+         module_as_package: bool,
+         setup_json: TextIO,
+         label: str,
+         proto_library: str = '',
+         proto_library_file: Path = None) -> int:
+    """Generates a setup.py and other files for a Python package."""
+    if proto_library_file:
+        try:
+            _check_nested_protos(label, proto_library_file, proto_library)
+        except ValueError as error:
+            msg = '\n'.join(textwrap.wrap(str(error), 78))
+            print(
+                f'ERROR: Failed to generate Python package {label}:\n\n'
+                f'{textwrap.indent(msg, "  ")}\n',
+                file=sys.stderr)
+            return 1
+
+    pkg_data = _collect_all_files(files, root, file_list, file_list_root)
+
+    if module_as_package:
+        _import_module_in_package_init(files)
+
+    # Create the setup.py file for this package.
+    root.joinpath('setup.py').write_text(
+        _generate_setup_py(pkg_data, setup_json))
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main(**vars(_parse_args())))
diff --git a/pw_build/python.gni b/pw_build/python.gni
index 9ecae0e..d04d027 100644
--- a/pw_build/python.gni
+++ b/pw_build/python.gni
@@ -15,10 +15,11 @@
 import("//build_overrides/pigweed.gni")
 
 import("$dir_pw_build/input_group.gni")
+import("$dir_pw_build/mirror_tree.gni")
 import("$dir_pw_build/python_action.gni")
 
 # Python packages provide the following targets as $target_name.$subtarget.
-_python_subtargets = [
+pw_python_package_subtargets = [
   "tests",
   "lint",
   "lint.mypy",
@@ -47,12 +48,16 @@
 # Args:
 #   setup: List of setup file paths (setup.py or pyproject.toml & setup.cfg),
 #       which must all be in the same directory.
+#   generate_setup: As an alternative to 'setup', generate setup files with the
+#       keywords in this scope. 'name' is required.
 #   sources: Python sources files in the package.
 #   tests: Test files for this Python package.
 #   python_deps: Dependencies on other pw_python_packages in the GN build.
 #   python_test_deps: Test-only pw_python_package dependencies.
 #   other_deps: Dependencies on GN targets that are not pw_python_packages.
 #   inputs: Other files to track, such as package_data.
+#   proto_library: A pw_proto_library target to embed in this Python package.
+#       generate_setup is required in place of setup if proto_library is used.
 #   lint: If true (default), applies mypy and pylint to the package. If false,
 #       does not.
 #   pylintrc: Optional path to a pylintrc configuration file to use. If not
@@ -64,18 +69,6 @@
 #       executed from the package's setup directory, so mypy.ini files in that
 #       directory will take precedence over others.
 template("pw_python_package") {
-  if (defined(invoker.sources)) {
-    _all_py_files = invoker.sources
-  } else {
-    _all_py_files = []
-  }
-
-  if (defined(invoker.tests)) {
-    _test_sources = invoker.tests
-  } else {
-    _test_sources = []
-  }
-
   # The Python targets are always instantiated in the default toolchain. Use
   # fully qualified labels so that the toolchain is not lost.
   _other_deps = []
@@ -85,64 +78,111 @@
     }
   }
 
-  _all_py_files += _test_sources
+  _python_deps = []
+  if (defined(invoker.python_deps)) {
+    foreach(dep, invoker.python_deps) {
+      _python_deps += [ get_label_info(dep, "label_with_toolchain") ]
+    }
+  }
 
   # pw_python_script uses pw_python_package, but with a limited set of features.
   # _pw_standalone signals that this target is actually a pw_python_script.
   _is_package = !(defined(invoker._pw_standalone) && invoker._pw_standalone)
 
-  # Some build targets generate Python packages, setting _pw_generated to
-  # indicate this.
-  _is_generated_package =
-      defined(invoker._pw_generated) && invoker._pw_generated
+  _generate_package = false
 
-  # Argument: invoker.lint = [true | false]; default = true.
-  # Default to false for generated packages, but allow overrides.
+  # Check the generate_setup and import_protos args to determine if this package
+  # is generated.
+  if (_is_package) {
+    assert(defined(invoker.generate_setup) != defined(invoker.setup),
+           "Either 'setup' or 'generate_setup' (but not both) must provided")
+
+    if (defined(invoker.proto_library)) {
+      assert(invoker.proto_library != "", "'proto_library' cannot be empty")
+      assert(defined(invoker.generate_setup),
+             "Python packages that import protos with 'proto_library' must " +
+                 "use 'generate_setup' instead of 'setup'")
+
+      _import_protos = [ invoker.proto_library ]
+    } else if (defined(invoker.generate_setup)) {
+      _import_protos = []
+    }
+
+    if (defined(invoker.generate_setup)) {
+      _generate_package = true
+      _setup_dir = "$target_gen_dir/$target_name.generated_python_package"
+
+      if (defined(invoker.strip_prefix)) {
+        _source_root = invoker.strip_prefix
+      } else {
+        _source_root = "."
+      }
+    } else {
+      # Non-generated packages with sources provided need an __init__.py.
+      assert(!defined(invoker.sources) || invoker.sources == [] ||
+                 filter_include(invoker.sources, [ "*\b__init__.py" ]) != [],
+             "Python packages must have at least one __init__.py file")
+
+      # Get the directories of the setup files. All must be in the same dir.
+      _setup_dirs = get_path_info(invoker.setup, "dir")
+      _setup_dir = _setup_dirs[0]
+
+      foreach(dir, _setup_dirs) {
+        assert(dir == _setup_dir,
+               "All files in 'setup' must be in the same directory")
+      }
+
+      assert(!defined(invoker.strip_prefix),
+             "'strip_prefix' may only be given if 'generate_setup' is provided")
+    }
+  }
+
+  # Process arguments defaults and set defaults.
+
+  # Argument: lint (bool); default = true.
   if (defined(invoker.lint)) {
     _should_lint = invoker.lint
   } else {
-    _should_lint = !_is_generated_package
+    _should_lint = true
   }
 
-  if (_is_package) {
-    assert(defined(invoker.setup) && invoker.setup != [],
-           "pw_python_package requires 'setup' to point to a setup.py file " +
-               "or pyproject.toml and setup.cfg files")
-
-    if (!_is_generated_package) {
-      _all_py_files += invoker.setup
-    }
-
-    # Get the directories of the setup files. All files must be in the same dir.
-    _setup_dirs = get_path_info(invoker.setup, "dir")
-    _setup_dir = _setup_dirs[0]
-
-    foreach(dir, _setup_dirs) {
-      assert(dir == _setup_dir,
-             "All files in 'setup' must be in the same directory")
-    }
-
-    # If sources are provided, make sure there is an __init__.py file.
-    if (!_is_generated_package && defined(invoker.sources) &&
-        invoker.sources != []) {
-      assert(filter_include(invoker.sources, [ "*\b__init__.py" ]) != [],
-             "Python packages must have at least one __init__.py file")
+  # Argument: sources (list)
+  _sources = []
+  if (defined(invoker.sources)) {
+    if (_generate_package) {
+      foreach(source, rebase_path(invoker.sources, _source_root)) {
+        _sources += [ "$_setup_dir/$source" ]
+      }
+    } else {
+      _sources += invoker.sources
     }
   }
 
-  _python_deps = []
-  if (defined(invoker.python_deps)) {
-    foreach(dep, invoker.python_deps) {
-      # Use the fully qualified name so the subtarget can be appended as needed.
-      _python_deps += [ get_label_info(dep, "label_no_toolchain") ]
+  # Argument: tests (list)
+  _test_sources = []
+  if (defined(invoker.tests)) {
+    if (_generate_package) {
+      foreach(source, rebase_path(invoker.tests, _source_root)) {
+        _test_sources += [ "$_setup_dir/$source" ]
+      }
+    } else {
+      _test_sources += invoker.tests
     }
   }
 
-  # All dependencies needed for the package and its tests.
-  _python_test_deps = _python_deps
+  # Argument: setup (list)
+  _setup_sources = []
+  if (defined(invoker.setup)) {
+    _setup_sources = invoker.setup
+  } else if (_generate_package) {
+    _setup_sources = [ "$_setup_dir/setup.py" ]
+  }
+
+  # Argument: python_test_deps (list)
+  _python_test_deps = _python_deps  # include all deps in test deps
   if (defined(invoker.python_test_deps)) {
-    foreach(test_dep, invoker.python_test_deps) {
-      _python_test_deps += [ get_label_info(test_dep, "label_no_toolchain") ]
+    foreach(dep, invoker.python_test_deps) {
+      _python_test_deps += [ get_label_info(dep, "label_with_toolchain") ]
     }
   }
 
@@ -150,220 +190,339 @@
     assert(!defined(invoker.python_test_deps),
            "python_test_deps was provided, but there are no tests in " +
                get_label_info(":$target_name", "label_no_toolchain"))
-    not_needed(_python_test_deps)
+    not_needed([ "_python_test_deps" ])
   }
 
-  _internal_target = "$target_name._internal"
+  _all_py_files = _sources + _test_sources + _setup_sources
 
-  # Create groups with the public target names ($target_name, $target_name.lint,
-  # $target_name.install, etc.). These are actually wrappers around internal
-  # Python actions instantiated with the default toolchain. This ensures there
-  # is only a single copy of each Python action in the build.
-  #
-  # The $target_name.tests group is created separately below.
-  foreach(subtarget, _python_subtargets - [ "tests" ]) {
-    group("$target_name.$subtarget") {
-      deps = [ ":$_internal_target.$subtarget($default_toolchain)" ]
-    }
-  }
+  # The pw_python_package subtargets are only instantiated in the default
+  # toolchain. Other toolchains just refer to targets in the default toolchain.
+  if (current_toolchain == default_toolchain) {
+    # Declare the main Python package group. This represents the Python files,
+    # but does not take any actions. GN targets can depend on the package name
+    # to run when any files in the package change.
+    if (_generate_package) {
+      # If this package is generated, mirror the sources to the final directory.
+      pw_mirror_tree("$target_name._mirror_sources_to_out_dir") {
+        directory = _setup_dir
 
-  group("$target_name") {
-    deps = [ ":$_internal_target($default_toolchain)" ]
-  }
+        sources = []
+        if (defined(invoker.sources)) {
+          sources += invoker.sources
+        }
+        if (defined(invoker.tests)) {
+          sources += invoker.tests
+        }
 
-  # Declare the main Python package group. This represents the Python files, but
-  # does not take any actions. GN targets can depend on the package name to run
-  # when any files in the package change.
-  pw_input_group("$_internal_target") {
-    inputs = _all_py_files
-    if (defined(invoker.inputs)) {
-      inputs += invoker.inputs
-    }
-
-    deps = _python_deps + _other_deps
-  }
-
-  if (_is_package) {
-    # Install this Python package and its dependencies in the current Python
-    # environment.
-    pw_python_action("$_internal_target.install") {
-      module = "pip"
-      args = [ "install" ]
-
-      # Don't install generated packages with --editable, since the build
-      # directory is ephemeral.
-      if (!_is_generated_package) {
-        args += [ "--editable" ]
+        source_root = _source_root
+        public_deps = _python_deps + _other_deps
       }
-      args += [ rebase_path(_setup_dir) ]
 
-      stamp = true
-
-      # Parallel pip installations don't work, so serialize pip invocations.
-      pool = "$dir_pw_build:pip_pool"
-
-      deps = [ ":$_internal_target" ]
-      foreach(dep, _python_deps) {
-        _subtarget = get_label_info(dep, "label_no_toolchain") + ".install"
-        deps += [ "$_subtarget(" + get_label_info(dep, "toolchain") + ")" ]
+      # Depend on the proto's _gen targets (from the default toolchain).
+      _gen_protos = []
+      foreach(proto, _import_protos) {
+        _gen_protos +=
+            [ get_label_info(proto, "label_no_toolchain") + ".python._gen" ]
       }
-    }
 
-    # TODO(pwbug/239): Add support for building groups of wheels. The code below
-    #     is incomplete and untested.
-    pw_python_action("$_internal_target.wheel") {
-      script = "$dir_pw_build/py/pw_build/python_wheels.py"
-
-      args = [
-        "--out_dir",
-        rebase_path(target_out_dir),
-      ]
-      args += rebase_path(_all_py_files)
-
-      deps = [ ":$_internal_target.install" ]
-      stamp = true
-    }
-  } else {
-    # If this is not a package, install or build wheels for its deps only.
-    group("$_internal_target.install") {
-      deps = []
-      foreach(dep, _python_deps) {
-        deps += [ "$dep.install" ]
+      generated_file("$target_name._protos") {
+        deps = _gen_protos
+        data_keys = [ "protoc_outputs" ]
+        outputs = [ "$_setup_dir/protos.txt" ]
       }
-    }
-    group("$_internal_target.wheel") {
-      deps = []
-      foreach(dep, _python_deps) {
-        deps += [ "$dep.wheel" ]
+
+      _protos_file = get_target_outputs(":${invoker.target_name}._protos")
+
+      generated_file("$target_name._protos_root") {
+        deps = _gen_protos
+        data_keys = [ "root" ]
+        outputs = [ "$_setup_dir/proto_root.txt" ]
       }
-    }
-  }
 
-  # Define the static analysis targets for this package.
-  group("$_internal_target.lint") {
-    deps = [
-      ":$_internal_target.lint.mypy",
-      ":$_internal_target.lint.pylint",
-    ]
-  }
+      _root_file = get_target_outputs(":${invoker.target_name}._protos_root")
 
-  if (_should_lint || _test_sources != []) {
-    # Packages that must be installed to use the package or run its tests.
-    _test_install_deps = [ ":$_internal_target.install" ]
-    foreach(dep, _python_test_deps + [ "$dir_pw_build:python_lint" ]) {
-      _test_install_deps += [ "$dep.install" ]
-    }
-  }
+      # Get generated_setup scope and write it to disk ask JSON.
+      _gen_setup = invoker.generate_setup
+      assert(defined(_gen_setup.name), "'name' is required in generate_package")
+      assert(!defined(_gen_setup.packages) && !defined(_gen_setup.package_data),
+             "'packages' and 'package_data' may not be provided " +
+                 "in 'generate_package'")
+      write_file("$_setup_dir/setup.json", _gen_setup, "json")
 
-  # For packages that are not generated, create targets to run mypy and pylint.
-  # Linting is not performed on generated packages.
-  if (_should_lint) {
-    # Run lint tools from the setup or target directory so that the tools detect
-    # config files (e.g. pylintrc or mypy.ini) in that directory. Config files
-    # may be explicitly specified with the pylintrc or mypy_ini arguments.
-    if (defined(_setup_dir)) {
-      _lint_directory = rebase_path(_setup_dir)
+      # Generate the setup.py, py.typed, and __init__.py files as needed.
+      action(target_name) {
+        script = "$dir_pw_build/py/pw_build/generate_python_package.py"
+        args = [
+                 "--label",
+                 get_label_info(":$target_name", "label_no_toolchain"),
+                 "--root",
+                 rebase_path(_setup_dir),
+                 "--setup-json",
+                 rebase_path("$_setup_dir/setup.json"),
+                 "--file-list",
+                 rebase_path(_protos_file[0]),
+                 "--file-list-root",
+                 rebase_path(_root_file[0]),
+               ] + rebase_path(_sources)
+
+        if (defined(invoker._pw_module_as_package) &&
+            invoker._pw_module_as_package) {
+          args += [ "--module-as-package" ]
+        }
+
+        public_deps = [
+          ":$target_name._mirror_sources_to_out_dir",
+          ":$target_name._protos",
+          ":$target_name._protos_root",
+        ]
+
+        foreach(proto, _import_protos) {
+          _tgt = get_label_info(proto, "label_no_toolchain")
+          _path = get_label_info("$_tgt($default_toolchain)", "target_gen_dir")
+          _name = get_label_info(_tgt, "name")
+
+          args += [
+            "--proto-library=$_tgt",
+            "--proto-library-file",
+            rebase_path("$_path/$_name.proto_library/python_package.txt"),
+          ]
+
+          public_deps += [ "$_tgt.python._gen($default_toolchain)" ]
+        }
+
+        outputs = _setup_sources
+      }
     } else {
-      _lint_directory = rebase_path(".")
-    }
+      # If the package is not generated, use an input group for the sources.
+      pw_input_group(target_name) {
+        inputs = _all_py_files
+        if (defined(invoker.inputs)) {
+          inputs += invoker.inputs
+        }
 
-    pw_python_action("$_internal_target.lint.mypy") {
-      module = "mypy"
-      args = [
-        "--pretty",
-        "--show-error-codes",
-      ]
-
-      if (defined(invoker.mypy_ini)) {
-        args += [ "--config-file=" + rebase_path(invoker.mypy_ini) ]
-        inputs = [ invoker.mypy_ini ]
-      }
-
-      args += rebase_path(_all_py_files)
-
-      # Use this environment variable to force mypy to colorize output.
-      # See https://github.com/python/mypy/issues/7771
-      environment = [ "MYPY_FORCE_COLOR=1" ]
-
-      directory = _lint_directory
-      stamp = true
-
-      deps = _test_install_deps
-
-      foreach(dep, _python_test_deps) {
-        deps += [ "$dep.lint.mypy" ]
+        deps = _python_deps + _other_deps
       }
     }
 
-    # Create a target to run pylint on each of the Python files in this
-    # package and its dependencies.
-    pw_python_action_foreach("$_internal_target.lint.pylint") {
-      module = "pylint"
-      args = [
-        rebase_path(".") + "/{{source_target_relative}}",
-        "--jobs=1",
-        "--output-format=colorized",
+    if (_is_package) {
+      # Install this Python package and its dependencies in the current Python
+      # environment.
+      pw_python_action("$target_name.install") {
+        module = "pip"
+        public_deps = []
+
+        args = [ "install" ]
+
+        # For generated packages, reinstall when any files change. For regular
+        # packages, only reinstall when setup.py changes.
+        if (_generate_package) {
+          public_deps += [ ":${invoker.target_name}" ]
+        } else {
+          inputs = invoker.setup
+
+          # Install with --editable since the complete package is in source.
+          args += [ "--editable" ]
+        }
+
+        args += [ rebase_path(_setup_dir) ]
+
+        stamp = true
+
+        # Parallel pip installations don't work, so serialize pip invocations.
+        pool = "$dir_pw_build:pip_pool"
+
+        foreach(dep, _python_deps) {
+          # We need to add a suffix to the target name, but the label is
+          # formatted as "//path/to:target(toolchain)", so we can't just append
+          # ".subtarget". Instead, we replace the opening parenthesis of the
+          # toolchain with ".suffix(".
+          public_deps += [ string_replace(dep, "(", ".install(") ]
+        }
+      }
+
+      # TODO(pwbug/239): Add support for building groups of wheels. The code below
+      #     is incomplete and untested.
+      pw_python_action("$target_name.wheel") {
+        script = "$dir_pw_build/py/pw_build/python_wheels.py"
+
+        args = [
+          "--out_dir",
+          rebase_path(target_out_dir),
+        ]
+        args += rebase_path(_all_py_files)
+
+        deps = [ ":${invoker.target_name}.install" ]
+        stamp = true
+      }
+    } else {
+      # If this is not a package, install or build wheels for its deps only.
+      group("$target_name.install") {
+        deps = []
+        foreach(dep, _python_deps) {
+          deps += [ string_replace(dep, "(", ".install(") ]
+        }
+      }
+      group("$target_name.wheel") {
+        deps = []
+        foreach(dep, _python_deps) {
+          deps += [ string_replace(dep, "(", ".wheel(") ]
+        }
+      }
+    }
+
+    # Define the static analysis targets for this package.
+    group("$target_name.lint") {
+      deps = [
+        ":${invoker.target_name}.lint.mypy",
+        ":${invoker.target_name}.lint.pylint",
       ]
+    }
 
-      if (defined(invoker.pylintrc)) {
-        args += [ "--rcfile=" + rebase_path(invoker.pylintrc) ]
-        inputs = [ invoker.pylintrc ]
+    if (_should_lint || _test_sources != []) {
+      # All packages to install for either general use or test running.
+      _test_install_deps = [ ":$target_name.install" ]
+      foreach(dep, _python_test_deps + [ "$dir_pw_build:python_lint" ]) {
+        _test_install_deps += [ string_replace(dep, "(", ".install(") ]
+      }
+    }
+
+    # For packages that are not generated, create targets to run mypy and pylint.
+    if (_should_lint) {
+      # Run lint tools from the setup or target directory so that the tools detect
+      # config files (e.g. pylintrc or mypy.ini) in that directory. Config files
+      # may be explicitly specified with the pylintrc or mypy_ini arguments.
+      if (defined(_setup_dir)) {
+        _lint_directory = rebase_path(_setup_dir)
+      } else {
+        _lint_directory = rebase_path(".")
       }
 
-      if (host_os == "win") {
-        # Allow CRLF on Windows, in case Git is set to switch line endings.
-        args += [ "--disable=unexpected-line-ending-format" ]
+      pw_python_action("$target_name.lint.mypy") {
+        module = "mypy"
+        args = [
+          "--pretty",
+          "--show-error-codes",
+        ]
+
+        if (defined(invoker.mypy_ini)) {
+          args += [ "--config-file=" + rebase_path(invoker.mypy_ini) ]
+          inputs = [ invoker.mypy_ini ]
+        }
+
+        args += rebase_path(_all_py_files)
+
+        # Use this environment variable to force mypy to colorize output.
+        # See https://github.com/python/mypy/issues/7771
+        environment = [ "MYPY_FORCE_COLOR=1" ]
+
+        directory = _lint_directory
+        stamp = true
+
+        deps = _test_install_deps
+
+        foreach(dep, _python_test_deps) {
+          deps += [ string_replace(dep, "(", ".lint.mypy(") ]
+        }
       }
 
-      sources = _all_py_files
+      # Create a target to run pylint on each of the Python files in this
+      # package and its dependencies.
+      pw_python_action_foreach("$target_name.lint.pylint") {
+        module = "pylint"
+        args = [
+          rebase_path(".") + "/{{source_target_relative}}",
+          "--jobs=1",
+          "--output-format=colorized",
+        ]
 
-      directory = _lint_directory
-      stamp = "$target_gen_dir/{{source_target_relative}}.pylint.passed"
+        if (defined(invoker.pylintrc)) {
+          args += [ "--rcfile=" + rebase_path(invoker.pylintrc) ]
+          inputs = [ invoker.pylintrc ]
+        }
 
-      deps = _test_install_deps
+        if (host_os == "win") {
+          # Allow CRLF on Windows, in case Git is set to switch line endings.
+          args += [ "--disable=unexpected-line-ending-format" ]
+        }
 
-      foreach(dep, _python_test_deps) {
-        deps += [ "$dep.lint.pylint" ]
+        sources = _all_py_files
+
+        directory = _lint_directory
+        stamp = "$target_gen_dir/{{source_target_relative}}.pylint.passed"
+
+        public_deps = _test_install_deps
+
+        foreach(dep, _python_test_deps) {
+          public_deps += [ string_replace(dep, "(", ".lint.pylint(") ]
+        }
+      }
+    } else {
+      pw_input_group("$target_name.lint.mypy") {
+        if (defined(invoker.pylintrc)) {
+          inputs = [ invoker.pylintrc ]
+        }
+      }
+      pw_input_group("$target_name.lint.pylint") {
+        if (defined(invoker.mypy_ini)) {
+          inputs = [ invoker.mypy_ini ]
+        }
       }
     }
   } else {
-    pw_input_group("$_internal_target.lint.mypy") {
-      if (defined(invoker.pylintrc)) {
-        inputs = [ invoker.pylintrc ]
+    # Create groups with the public target names ($target_name, $target_name.lint,
+    # $target_name.install, etc.). These are actually wrappers around internal
+    # Python actions instantiated with the default toolchain. This ensures there
+    # is only a single copy of each Python action in the build.
+    #
+    # The $target_name.tests group is created separately below.
+    group("$target_name") {
+      deps = [ ":$target_name($default_toolchain)" ]
+    }
+
+    foreach(subtarget, pw_python_package_subtargets - [ "tests" ]) {
+      group("$target_name.$subtarget") {
+        deps = [ ":${invoker.target_name}.$subtarget($default_toolchain)" ]
       }
     }
-    pw_input_group("$_internal_target.lint.pylint") {
-      if (defined(invoker.mypy_ini)) {
-        inputs = [ invoker.mypy_ini ]
-      }
-    }
+
+    # Everything Python-related is only instantiated in the default toolchain.
+    # Silence not-needed warnings except for in the default toolchain.
+    not_needed("*")
+    not_needed(invoker, "*")
   }
 
   # Create a target for each test file.
   _test_targets = []
 
   foreach(test, _test_sources) {
-    _test_name = string_replace(test, "/", "_")
-    _internal_test_target = "$_internal_target.tests.$_test_name"
+    if (_is_package) {
+      _name = rebase_path(test, _setup_dir)
+    } else {
+      _name = test
+    }
 
-    pw_python_action(_internal_test_target) {
-      script = test
-      stamp = true
+    _test_target = "$target_name.tests." + string_replace(_name, "/", "_")
 
-      deps = _test_install_deps
+    if (current_toolchain == default_toolchain) {
+      pw_python_action(_test_target) {
+        script = test
+        stamp = true
 
-      foreach(dep, _python_test_deps) {
-        deps += [ "$dep.tests" ]
+        deps = _test_install_deps
+
+        foreach(dep, _python_test_deps) {
+          deps += [ string_replace(dep, "(", ".tests(") ]
+        }
+      }
+    } else {
+      # Create a public version of each test target, so tests can be executed as
+      # //path/to:package.tests.foo.py.
+      group(_test_target) {
+        deps = [ ":$_test_target($default_toolchain)" ]
       }
     }
 
-    # Create a public version of each test target, so tests can be executed as
-    # //path/to:package.tests.foo.py.
-    group("$target_name.tests.$_test_name") {
-      deps = [ ":$_internal_test_target" ]
-    }
-
-    _test_targets += [ ":$target_name.tests.$_test_name" ]
+    _test_targets += [ ":$_test_target" ]
   }
 
   group("$target_name.tests") {
@@ -386,7 +545,7 @@
     deps = _python_deps
   }
 
-  foreach(subtarget, _python_subtargets) {
+  foreach(subtarget, pw_python_package_subtargets) {
     group("$target_name.$subtarget") {
       deps = []
       foreach(dep, _python_deps) {
@@ -472,7 +631,7 @@
 
   # Create stubs for the unused subtargets so that pw_python_requirements can be
   # used as python_deps.
-  foreach(subtarget, _python_subtargets - [ "install" ]) {
+  foreach(subtarget, pw_python_package_subtargets - [ "install" ]) {
     group("$target_name.$subtarget") {
     }
   }
diff --git a/pw_build/python.rst b/pw_build/python.rst
index 11338cd..7157746 100644
--- a/pw_build/python.rst
+++ b/pw_build/python.rst
@@ -31,6 +31,9 @@
 - ``python_test_deps`` - Test-only pw_python_package dependencies.
 - ``other_deps`` - Dependencies on GN targets that are not pw_python_packages.
 - ``inputs`` - Other files to track, such as package_data.
+- ``proto_library`` - A pw_proto_library target to embed in this Python package.
+  generate_setup is required in place of setup if proto_library is used. See
+  :ref:`module-pw_protobuf_compiler-add-to-python-package`.
 - ``lint`` - If true (default), applies Mypy and Pylint to the package. If
   false, does not.
 - ``pylintrc`` - Optional path to a pylintrc configuration file to use. If not
diff --git a/pw_protobuf_compiler/docs.rst b/pw_protobuf_compiler/docs.rst
index ac659fb..8651c13 100644
--- a/pw_protobuf_compiler/docs.rst
+++ b/pw_protobuf_compiler/docs.rst
@@ -66,6 +66,8 @@
   compiled with protoc as ``"nested/foo.proto"``.
 * ``strip_prefix``: Remove this prefix from the source protos. All source and
   input files must be nested under this path.
+* ``python_package``: Label of Python package in which to nest the proto
+  modules.
 
 **Example**
 
@@ -149,6 +151,44 @@
   └── internal
       └── gamma.proto
 
+.. _module-pw_protobuf_compiler-add-to-python-package:
+
+Adding Python proto modules to an existing package
+--------------------------------------------------
+By default, generated Python proto modules are organized into their own Python
+package. These proto modules can instead be added to an existing Python package
+declared with ``pw_python_library``. This is done by setting the
+``python_package`` argument on the ``pw_proto_library`` and the
+``proto_library`` argument on the ``pw_python_package``.
+
+For example, the protos declared in ``my_protos`` will be nested in the Python
+package declared by ``my_package``.
+
+.. code-block::
+
+  pw_proto_library("my_protos") {
+    sources = [ "hello.proto ]
+    prefix = "foo"
+    python_package = ":my_package"
+  }
+
+  pw_python_pacakge("my_package") {
+    generate_setup = {
+      name = "foo"
+      version = "1.0"
+    }
+    sources = [ "foo/cool_module.py" ]
+    proto_library = ":my_protos"
+  }
+
+The ``hello_pb2.py`` proto module can be used alongside other files in the
+``foo`` package.
+
+.. code-block:: python
+
+  from foo import cool_module, hello_pb2
+
+
 Working with externally defined protos
 --------------------------------------
 ``pw_proto_library`` targets may be used to build ``.proto`` sources from
diff --git a/pw_protobuf_compiler/proto.gni b/pw_protobuf_compiler/proto.gni
index 27c608f..c877b28 100644
--- a/pw_protobuf_compiler/proto.gni
+++ b/pw_protobuf_compiler/proto.gni
@@ -49,7 +49,6 @@
       rebase_path(get_target_outputs(":${invoker.base_target}._includes"))
 
   pw_python_action("$target_name._gen") {
-    forward_variables_from(invoker, [ "metadata" ])
     script =
         "$dir_pw_protobuf_compiler/py/pw_protobuf_compiler/generate_protos.py"
 
@@ -89,6 +88,15 @@
     } else {
       stamp = true
     }
+
+    if (defined(invoker.metadata)) {
+      metadata = invoker.metadata
+    } else {
+      metadata = {
+        protoc_outputs = rebase_path(outputs)
+        root = [ rebase_path(_out_dir) ]
+      }
+    }
   }
 }
 
@@ -252,42 +260,49 @@
 # the generated files. This is internal and should not be used outside of this
 # file. Use pw_proto_library instead.
 template("_pw_python_proto_library") {
-  _target = target_name
-
   _pw_invoke_protoc(target_name) {
-    forward_variables_from(invoker, "*", _forwarded_vars)
+    forward_variables_from(invoker, "*", _forwarded_vars + [ "python_package" ])
     language = "python"
     python_deps = [ "$dir_pw_protobuf_compiler:protobuf_requirements" ]
   }
 
-  _setup_py = "${invoker.base_out_dir}/python/setup.py"
+  if (defined(invoker.python_package) && invoker.python_package != "") {
+    # If nested in a Python package, write the package's name to a file so
+    # pw_python_package can check that the dependencies are correct.
+    write_file("${invoker.base_out_dir}/python_package.txt",
+               get_label_info(invoker.python_package, "label_no_toolchain"))
 
-  # Create the setup and init files for the Python package.
-  action(target_name + "._package_gen") {
-    script = "$dir_pw_protobuf_compiler/py/pw_protobuf_compiler/generate_python_package.py"
-    args = [
-             "--setup",
-             rebase_path(_setup_py),
-             "--package",
-             invoker._package_dir,
-           ] + rebase_path(invoker.outputs, "${invoker.base_out_dir}/python")
-
-    if (invoker.module_as_package != "") {
-      args += [ "--module-as-package" ]
+    # If anyone attempts to depend on this Python package, print an error.
+    pw_error(target_name) {
+      _pkg = get_label_info(invoker.python_package, "label_no_toolchain")
+      message_lines = [
+        "This proto Python package is embedded in the $_pkg Python package.",
+        "It cannot be used directly; instead, depend on $_pkg.",
+      ]
     }
+    foreach(subtarget, pw_python_package_subtargets) {
+      group("$target_name.$subtarget") {
+        deps = [ ":${invoker.target_name}" ]
+      }
+    }
+  } else {
+    write_file("${invoker.base_out_dir}/python_package.txt", "")
 
-    public_deps = [ ":$_target._gen($default_toolchain)" ]
-    outputs = [ _setup_py ]
-  }
+    # Create a Python package with the generated source files.
+    pw_python_package(target_name) {
+      forward_variables_from(invoker, _forwarded_vars)
+      generate_setup = {
+        name = invoker._package_dir
+        version = "0.0.1"  # TODO(hepler): Need to be able to set this verison.
+      }
+      sources = invoker.outputs
+      strip_prefix = "${invoker.base_out_dir}/python"
+      python_deps = invoker.deps
+      other_deps = [ ":$target_name._gen($default_toolchain)" ]
+      lint = false
 
-  # Create a Python package with the generated source files.
-  pw_python_package(target_name) {
-    forward_variables_from(invoker, _forwarded_vars)
-    setup = [ _setup_py ]
-    sources = invoker.outputs
-    python_deps = invoker.deps
-    other_deps = [ ":$_target._package_gen($default_toolchain)" ]
-    _pw_generated = true
+      _pw_module_as_package = invoker.module_as_package != ""
+    }
   }
 }
 
@@ -309,6 +324,7 @@
 #       compiled with protoc as "nested/foo.proto".
 #   strip_prefix: Remove this prefix from the source protos. All source and
 #       input files must be nested under this path.
+#   python_package: Label of Python package in which to nest the proto modules.
 #
 template("pw_proto_library") {
   assert(defined(invoker.sources) && invoker.sources != [],
@@ -348,8 +364,9 @@
     # This is the output directory for all files related to this proto library.
     # Sources are mirrored to "$base_out_dir/sources" and protoc puts outputs in
     # "$base_out_dir/$language" by default.
-    base_out_dir = get_label_info(":$target_name($default_toolchain)",
-                                  "target_gen_dir") + "/$target_name"
+    base_out_dir =
+        get_label_info(":$target_name($default_toolchain)", "target_gen_dir") +
+        "/$target_name.proto_library"
 
     compile_dir = "$base_out_dir/sources"
 
@@ -422,7 +439,7 @@
     deps = process_file_template(_deps, "{{source}}._includes")
 
     data_keys = [ "protoc_includes" ]
-    outputs = [ "$target_gen_dir/${_common.base_target}/includes.txt" ]
+    outputs = [ "${_common.base_out_dir}/includes.txt" ]
 
     # Indicate this library's base directory for its dependents.
     metadata = {
@@ -537,6 +554,7 @@
 
   _pw_python_proto_library("$target_name.python") {
     forward_variables_from(_common, "*")
+    forward_variables_from(invoker, [ "python_package" ])
     module_as_package = _module_as_package
 
     deps = []
diff --git a/pw_protobuf_compiler/py/BUILD.gn b/pw_protobuf_compiler/py/BUILD.gn
index 1d5a10b..687cb1a 100644
--- a/pw_protobuf_compiler/py/BUILD.gn
+++ b/pw_protobuf_compiler/py/BUILD.gn
@@ -22,7 +22,6 @@
   sources = [
     "pw_protobuf_compiler/__init__.py",
     "pw_protobuf_compiler/generate_protos.py",
-    "pw_protobuf_compiler/generate_python_package.py",
     "pw_protobuf_compiler/proto_target_invalid.py",
     "pw_protobuf_compiler/python_protos.py",
   ]
diff --git a/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_python_package.py b/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_python_package.py
deleted file mode 100644
index 9147f08..0000000
--- a/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_python_package.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# Copyright 2020 The Pigweed Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License"); you may not
-# use this file except in compliance with the License. You may obtain a copy of
-# the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
-# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
-# License for the specific language governing permissions and limitations under
-# the License.
-"""Generates a setup.py and __init__.py for a Python package."""
-
-import argparse
-from collections import defaultdict
-from pathlib import Path
-import sys
-from typing import Dict, List, Set
-
-# Make sure dependencies are optional, since this script may be run when
-# installing Python package dependencies through GN.
-try:
-    from pw_cli.log import install as setup_logging
-except ImportError:
-    from logging import basicConfig as setup_logging  # type: ignore
-
-_SETUP_TEMPLATE = """# Generated file. Do not modify.
-import setuptools
-
-setuptools.setup(
-    name={name!r},
-    version='0.0.1',
-    author='Pigweed Authors',
-    author_email='pigweed-developers@googlegroups.com',
-    description='Generated protobuf files',
-    packages={packages!r},
-    package_data={package_data!r},
-    include_package_data=True,
-    zip_safe=False,
-    install_requires=['protobuf'],
-)
-"""
-
-
-def _parse_args():
-    """Parses and returns the command line arguments."""
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('--package',
-                        required=True,
-                        help='Name of the generated Python package')
-    parser.add_argument('--setup',
-                        required=True,
-                        type=Path,
-                        help='Path to setup.py file')
-    parser.add_argument('--module-as-package',
-                        action='store_true',
-                        help='The package is a standalone external proto')
-    parser.add_argument('sources',
-                        type=Path,
-                        nargs='+',
-                        help='Relative paths to the .py and .pyi files')
-    return parser.parse_args()
-
-
-def main(package: str, setup: Path, module_as_package: bool,
-         sources: List[Path]) -> int:
-    """Generates __init__.py and py.typed files and a setup.py."""
-    assert not module_as_package or len(sources) == 2
-
-    base = setup.parent.resolve()
-    base.mkdir(exist_ok=True)
-
-    # Find all directories in the package, including empty ones.
-    subpackages: Set[Path] = set()
-    for source in sources:
-        subpackages.update(base / path for path in source.parents)
-    subpackages.remove(base)
-
-    pkg_data: Dict[str, List[str]] = defaultdict(list)
-
-    # Create __init__.py and py.typed files for each subdirectory.
-    for pkg in subpackages:
-        pkg.mkdir(exist_ok=True, parents=True)
-        pkg.joinpath('__init__.py').write_text('')
-
-        package_name = pkg.relative_to(base).as_posix().replace('/', '.')
-        pkg.joinpath('py.typed').touch()
-        pkg_data[package_name].append('py.typed')
-
-    # Add the Mypy stub (.pyi) for each source file.
-    for mypy_stub in (s for s in sources if s.suffix == '.pyi'):
-        pkg = base / mypy_stub.parent
-        package_name = pkg.relative_to(base).as_posix().replace('/', '.')
-        pkg_data[package_name].append(mypy_stub.name)
-
-        if module_as_package:
-            pkg.joinpath('__init__.py').write_text(
-                f'from {mypy_stub.stem}.{mypy_stub.stem} import *\n')
-
-    setup.write_text(
-        _SETUP_TEMPLATE.format(name=package,
-                               packages=list(pkg_data),
-                               package_data=dict(pkg_data)))
-
-    return 0
-
-
-if __name__ == '__main__':
-    setup_logging()
-    sys.exit(main(**vars(_parse_args())))