third_party/pigweed/repo/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_python_package.py

   1 # Copyright 2020 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """Generates a setup.py and __init__.py for a Python package."""
  15
  16 import argparse
  17 from collections import defaultdict
  18 from pathlib import Path
  19 import sys
  20 from typing import Dict, List, Set
  21
  22 # Make sure dependencies are optional, since this script may be run when
  23 # installing Python package dependencies through GN.
  24 try:
  25     from pw_cli.log import install as setup_logging
  26 except ImportError:
  27     from logging import basicConfig as setup_logging  # type: ignore
  28
  29 _SETUP_TEMPLATE = """# Generated file. Do not modify.
  30 import setuptools
  31
  32 setuptools.setup(
  33     name={name!r},
  34     version='0.0.1',
  35     author='Pigweed Authors',
  36     author_email='pigweed-developers@googlegroups.com',
  37     description='Generated protobuf files',
  38     packages={packages!r},
  39     package_data={package_data!r},
  40     include_package_data=True,
  41     zip_safe=False,
  42     install_requires=['protobuf'],
  43 )
  44 """
  45
  46
  47 def _parse_args():
  48     """Parses and returns the command line arguments."""
  49     parser = argparse.ArgumentParser(description=__doc__)
  50     parser.add_argument('--package',
  51                         required=True,
  52                         help='Name of the generated Python package')
  53     parser.add_argument('--setup',
  54                         required=True,
  55                         type=Path,
  56                         help='Path to setup.py file')
  57     parser.add_argument('--standalone',
  58                         action='store_true',
  59                         help='The package is a standalone external proto')
  60     parser.add_argument('sources',
  61                         type=Path,
  62                         nargs='+',
  63                         help='Relative paths to the .py and .pyi files')
  64     return parser.parse_args()
  65
  66
  67 def main(package: str, setup: Path, standalone: bool,
  68          sources: List[Path]) -> int:
  69     """Generates __init__.py and py.typed files and a setup.py."""
  70     assert not standalone or len(sources) == 2
  71
  72     base = setup.parent.resolve()
  73     base.mkdir(exist_ok=True)
  74
  75     # Find all directories in the package, including empty ones.
  76     subpackages: Set[Path] = set()
  77     for source in sources:
  78         subpackages.update(base / path for path in source.parents)
  79     subpackages.remove(base)
  80
  81     pkg_data: Dict[str, List[str]] = defaultdict(list)
  82
  83     # Create __init__.py and py.typed files for each subdirectory.
  84     for pkg in subpackages:
  85         pkg.mkdir(exist_ok=True, parents=True)
  86         pkg.joinpath('__init__.py').write_text('')
  87
  88         package_name = pkg.relative_to(base).as_posix().replace('/', '.')
  89         pkg.joinpath('py.typed').touch()
  90         pkg_data[package_name].append('py.typed')
  91
  92     # Add the Mypy stub (.pyi) for each source file.
  93     for mypy_stub in (s for s in sources if s.suffix == '.pyi'):
  94         pkg = base / mypy_stub.parent
  95         package_name = pkg.relative_to(base).as_posix().replace('/', '.')
  96         pkg_data[package_name].append(mypy_stub.name)
  97
  98         if standalone:
  99             pkg.joinpath('__init__.py').write_text(
 100                 f'from {mypy_stub.stem}.{mypy_stub.stem} import *\n')
 101
 102     setup.write_text(
 103         _SETUP_TEMPLATE.format(name=package,
 104                                packages=list(pkg_data),
 105                                package_data=dict(pkg_data)))
 106
 107     return 0
 108
 109
 110 if __name__ == '__main__':
 111     setup_logging()
 112     sys.exit(main(**vars(_parse_args())))