pw_build: Update pw_proto_library metadata files

- Have pw_proto_library targets output a JSON file with information
  about the proto files. This consolidates the two files that were used
  previously into one and makes it possible to support dependencies in
  protos imported into Python packages.
- Allow other_deps in pw_python_group.

Change-Id: I1cd17a344a30e47efee371471881f350bb25ca41
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/45303
Reviewed-by: Keir Mierle <keir@google.com>
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Wyatt Hepler <hepler@google.com>
diff --git a/pw_build/py/pw_build/generate_python_package.py b/pw_build/py/pw_build/generate_python_package.py
index 81e6c48..e70ce71 100644
--- a/pw_build/py/pw_build/generate_python_package.py
+++ b/pw_build/py/pw_build/generate_python_package.py
@@ -15,11 +15,13 @@
 
 import argparse
 from collections import defaultdict
+from dataclasses import dataclass
+from itertools import chain
 import json
 from pathlib import Path
 import sys
 import textwrap
-from typing import Dict, List, Set, TextIO
+from typing import Any, Dict, Iterable, Iterator, List, Set, Sequence, TextIO
 
 try:
     from pw_build.mirror_tree import mirror_paths
@@ -33,22 +35,14 @@
 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',
+                        dest='proto_libraries',
+                        type=argparse.FileType('r'),
+                        default=[],
+                        action='append',
+                        help='Paths')
+    parser.add_argument('--generated-root',
                         required=True,
                         type=Path,
                         help='The base directory for the Python package')
@@ -66,35 +60,33 @@
     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'
+def _check_nested_protos(label: str, proto_info: Dict[str, Any]) -> None:
+    """Checks that the proto library refers to this package."""
+    python_package = proto_info['nested_in_python_package']
 
-    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}.")
+    if python_package != label:
+        raise ValueError(
+            f"{label}'s 'proto_library' is set to {proto_info['label']}, but "
+            f"that target's 'python_package' is {python_package or 'not set'}. "
+            f"Set {proto_info['label']}'s 'python_package' to {label}.")
 
 
-def _collect_all_files(files: List[Path], root: Path, file_list: TextIO,
-                       file_list_root: TextIO) -> Dict[str, Set[str]]:
+@dataclass(frozen=True)
+class _ProtoInfo:
+    root: Path
+    sources: Sequence[Path]
+    deps: Sequence[str]
+
+
+def _collect_all_files(
+        root: Path, files: List[Path],
+        paths_to_collect: Iterable[_ProtoInfo]) -> 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)
+    for proto_info in paths_to_collect:
+        # Mirror the proto files to this package.
+        files += mirror_paths(proto_info.root, proto_info.sources, root)
 
     # Find all subpackages, including empty ones.
     subpackages: Set[Path] = set()
@@ -132,18 +124,16 @@
 '''
 
 
-def _generate_setup_py(pkg_data: dict, setup_json: TextIO) -> str:
+def _generate_setup_py(pkg_data: dict, keywords: dict) -> 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), (
+    assert not any(kw in keywords for kw in setup_keywords), (
         'Generated packages may not specify "packages" or "package_data"')
-    setup_keywords.update(specified_keywords)
+    setup_keywords.update(keywords)
 
     return _SETUP_PY_FILE.format(keywords='\n'.join(
         f'    {k}={v!r},' for k, v in setup_keywords.items()))
@@ -166,35 +156,48 @@
         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
+def _load_metadata(label: str,
+                   proto_libraries: Iterable[TextIO]) -> Iterator[_ProtoInfo]:
+    for proto_library_file in proto_libraries:
+        info = json.load(proto_library_file)
+        _check_nested_protos(label, info)
 
-    pkg_data = _collect_all_files(files, root, file_list, file_list_root)
+        deps = []
+        for dep in info['dependencies']:
+            with open(dep) as file:
+                deps.append(json.load(file)['package'])
+
+        yield _ProtoInfo(Path(info['root']),
+                         tuple(Path(p) for p in info['protoc_outputs']), deps)
+
+
+def main(generated_root: Path, files: List[Path], module_as_package: bool,
+         setup_json: TextIO, label: str,
+         proto_libraries: Iterable[TextIO]) -> int:
+    """Generates a setup.py and other files for a Python package."""
+    proto_infos = list(_load_metadata(label, proto_libraries))
+    try:
+        pkg_data = _collect_all_files(generated_root, files, proto_infos)
+    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
+
+    with setup_json:
+        setup_keywords = json.load(setup_json)
+
+    install_requires = setup_keywords.setdefault('install_requires', [])
+    install_requires += chain.from_iterable(i.deps for i in proto_infos)
 
     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))
+    generated_root.joinpath('setup.py').write_text(
+        _generate_setup_py(pkg_data, setup_keywords))
 
     return 0
 
diff --git a/pw_build/python.gni b/pw_build/python.gni
index de56f91..4d4b014 100644
--- a/pw_build/python.gni
+++ b/pw_build/python.gni
@@ -196,6 +196,11 @@
                  "use 'generate_setup' instead of 'setup'")
 
       _import_protos = [ invoker.proto_library ]
+
+      # Depend on the dependencies of the proto library.
+      _proto = get_label_info(invoker.proto_library, "label_no_toolchain")
+      _toolchain = get_label_info(invoker.proto_library, "toolchain")
+      _python_deps += [ "$_proto.python._deps($_toolchain)" ]
     } else if (defined(invoker.generate_setup)) {
       _import_protos = []
     }
@@ -323,29 +328,6 @@
         public_deps = _python_deps + _other_deps
       }
 
-      # 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($pw_protobuf_compiler_TOOLCHAIN)" ]
-      }
-
-      generated_file("$target_name._protos") {
-        deps = _gen_protos
-        data_keys = [ "protoc_outputs" ]
-        outputs = [ "$_setup_dir/protos.txt" ]
-      }
-
-      _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" ]
-      }
-
-      _root_file = get_target_outputs(":${invoker.target_name}._protos_root")
-
       # 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")
@@ -360,16 +342,24 @@
         args = [
                  "--label",
                  get_label_info(":$target_name", "label_no_toolchain"),
-                 "--root",
+                 "--generated-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)
 
+        # Pass in the .json information files for the imported proto libraries.
+        foreach(proto, _import_protos) {
+          _label = get_label_info(proto, "label_no_toolchain") +
+                   ".python($pw_protobuf_compiler_TOOLCHAIN)"
+          _file = get_label_info(_label, "target_gen_dir") + "/" +
+                  get_label_info(_label, "name") + ".json"
+          args += [
+            "--proto-library",
+            rebase_path(_file),
+          ]
+        }
+
         if (defined(invoker._pw_module_as_package) &&
             invoker._pw_module_as_package) {
           args += [ "--module-as-package" ]
@@ -377,29 +367,7 @@
 
         inputs = [ "$_setup_dir/setup.json" ]
 
-        public_deps = [
-          ":$target_name._mirror_sources_to_out_dir",
-          ":$target_name._protos",
-          ":$target_name._protos_root",
-        ]
-
-        # Each pw_proto_library generates a file that indicates which Python
-        # package it is nested in, if any. Locate those files.
-        foreach(proto, _import_protos) {
-          _tgt = get_label_info(proto, "label_no_toolchain")
-          _path = get_label_info("$_tgt($pw_protobuf_compiler_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($pw_protobuf_compiler_TOOLCHAIN)" ]
-        }
+        public_deps = [ ":$target_name._mirror_sources_to_out_dir" ]
 
         outputs = _setup_sources
       }
@@ -647,6 +615,10 @@
 
   group(target_name) {
     deps = _python_deps
+
+    if (defined(invoker.other_deps)) {
+      deps += invoker.other_deps
+    }
   }
 
   foreach(subtarget, pw_python_package_subtargets) {
diff --git a/pw_protobuf_compiler/proto.gni b/pw_protobuf_compiler/proto.gni
index bc5622c..07f4d52 100644
--- a/pw_protobuf_compiler/proto.gni
+++ b/pw_protobuf_compiler/proto.gni
@@ -98,13 +98,30 @@
 
       if (defined(invoker.metadata)) {
         metadata = invoker.metadata
-      } else {
-        metadata = {
-          protoc_outputs = rebase_path(outputs)
-          root = [ rebase_path(_out_dir) ]
-        }
       }
     }
+
+    # Output a .json file with information about this proto library.
+    _proto_info = {
+      label = get_label_info(":${invoker.target_name}", "label_no_toolchain")
+      protoc_outputs = rebase_path(get_target_outputs(":$target_name._gen"))
+      root = rebase_path(_out_dir)
+      package = invoker.package
+
+      nested_in_python_package = ""
+      if (defined(invoker.python_package)) {
+        nested_in_python_package =
+            get_label_info(invoker.python_package, "label_no_toolchain")
+      }
+
+      dependencies = []
+      foreach(dep, invoker.deps) {
+        dependencies +=
+            rebase_path([ get_label_info(dep, "target_gen_dir") + "/" +
+                              get_label_info(dep, "name") + ".json" ])
+      }
+    }
+    write_file("$target_gen_dir/$target_name.json", _proto_info, "json")
   } else {
     # protoc is only ever invoked from pw_protobuf_compiler_TOOLCHAIN.
     not_needed([ "target_name" ])
@@ -284,14 +301,13 @@
     forward_variables_from(invoker, "*", _forwarded_vars + [ "python_package" ])
     language = "python"
     python_deps = [ "$dir_pw_protobuf_compiler:protobuf_requirements" ]
+
+    if (defined(invoker.python_package)) {
+      python_package = invoker.python_package
+    }
   }
 
   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"))
-
     # 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")
@@ -305,9 +321,15 @@
         deps = [ ":${invoker.target_name}" ]
       }
     }
-  } else {
-    write_file("${invoker.base_out_dir}/python_package.txt", "")
 
+    # This proto library is merged into another package, but create a target to
+    # collect its dependencies that the other package can depend on.
+    pw_python_group(target_name + "._deps") {
+      python_deps = invoker.deps
+      other_deps =
+          [ ":${invoker.target_name}._gen($pw_protobuf_compiler_TOOLCHAIN)" ]
+    }
+  } else {
     # Create a Python package with the generated source files.
     pw_python_package(target_name) {
       forward_variables_from(invoker, _forwarded_vars)
@@ -386,25 +408,6 @@
     _prefix = ""
   }
 
-  _common = {
-    base_target = target_name
-
-    # 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($pw_protobuf_compiler_TOOLCHAIN)",
-                       "target_gen_dir") + "/$target_name.proto_library"
-
-    compile_dir = "$base_out_dir/sources"
-
-    # Refer to the source files as the are mirrored to the output directory.
-    sources = []
-    foreach(file, rebase_path(invoker.sources, _source_root)) {
-      sources += [ "$compile_dir/$_prefix/$file" ]
-    }
-  }
-
   _package_dir = ""
   _source_names = []
 
@@ -460,6 +463,27 @@
     _deps = []
   }
 
+  _common = {
+    base_target = target_name
+
+    # 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($pw_protobuf_compiler_TOOLCHAIN)",
+                       "target_gen_dir") + "/$target_name.proto_library"
+
+    compile_dir = "$base_out_dir/sources"
+
+    # Refer to the source files as the are mirrored to the output directory.
+    sources = []
+    foreach(file, rebase_path(invoker.sources, _source_root)) {
+      sources += [ "$compile_dir/$_prefix/$file" ]
+    }
+
+    package = _package_dir
+  }
+
   # For each proto target, create a file which collects the base directories of
   # all of its dependencies to list as include paths to protoc.
   generated_file("$target_name._includes") {