|
| 1 | +"""Script to generate keras public API in `keras_cv/api` directory. |
| 2 | +
|
| 3 | +Usage: |
| 4 | +
|
| 5 | +Run via `./shell/api_gen.sh`. |
| 6 | +It generates API and formats user and generated APIs. |
| 7 | +""" |
| 8 | + |
| 9 | +import os |
| 10 | +import shutil |
| 11 | + |
| 12 | +import namex |
| 13 | + |
| 14 | +package = "keras_cv" |
| 15 | + |
| 16 | + |
| 17 | +def ignore_files(_, filenames): |
| 18 | + return [f for f in filenames if f.endswith("_test.py")] |
| 19 | + |
| 20 | + |
| 21 | +def copy_source_to_build_directory(root_path): |
| 22 | + # Copy sources (`keras_cv/` directory and setup files) to build dir |
| 23 | + build_dir = os.path.join(root_path, "tmp_build_dir") |
| 24 | + if os.path.exists(build_dir): |
| 25 | + shutil.rmtree(build_dir) |
| 26 | + os.mkdir(build_dir) |
| 27 | + shutil.copytree( |
| 28 | + package, os.path.join(build_dir, package), ignore=ignore_files |
| 29 | + ) |
| 30 | + return build_dir |
| 31 | + |
| 32 | + |
| 33 | +def export_version_string(api_init_fname): |
| 34 | + with open(api_init_fname) as f: |
| 35 | + contents = f.read() |
| 36 | + with open(api_init_fname, "w") as f: |
| 37 | + contents += "from keras_cv.src.version_utils import __version__\n" |
| 38 | + f.write(contents) |
| 39 | + |
| 40 | + |
| 41 | +def update_package_init(init_fname): |
| 42 | + contents = """ |
| 43 | +# Import everything from /api/ into keras. |
| 44 | +from keras_cv.api import * # noqa: F403 |
| 45 | +from keras_cv.api import __version__ # Import * ignores names start with "_". |
| 46 | +
|
| 47 | +import os |
| 48 | +
|
| 49 | +# Add everything in /api/ to the module search path. |
| 50 | +__path__.append(os.path.join(os.path.dirname(__file__), "api")) # noqa: F405 |
| 51 | +
|
| 52 | +# Don't pollute namespace. |
| 53 | +del os |
| 54 | +
|
| 55 | +# Never autocomplete `.src` or `.api` on an imported keras object. |
| 56 | +def __dir__(): |
| 57 | + keys = dict.fromkeys((globals().keys())) |
| 58 | + keys.pop("src") |
| 59 | + keys.pop("api") |
| 60 | + return list(keys) |
| 61 | +
|
| 62 | +
|
| 63 | +# Don't import `.src` or `.api` during `from keras import *`. |
| 64 | +__all__ = [ |
| 65 | + name |
| 66 | + for name in globals().keys() |
| 67 | + if not (name.startswith("_") or name in ("src", "api")) |
| 68 | +]""" |
| 69 | + with open(init_fname) as f: |
| 70 | + init_contents = f.read() |
| 71 | + with open(init_fname, "w") as f: |
| 72 | + f.write(init_contents.replace("\nfrom keras_cv import api", contents)) |
| 73 | + |
| 74 | + |
| 75 | +def build(): |
| 76 | + # Backup the `keras_cv/__init__.py` and restore it on error in api gen. |
| 77 | + root_path = os.path.dirname(os.path.abspath(__file__)) |
| 78 | + code_api_dir = os.path.join(root_path, package, "api") |
| 79 | + code_init_fname = os.path.join(root_path, package, "__init__.py") |
| 80 | + # Create temp build dir |
| 81 | + build_dir = copy_source_to_build_directory(root_path) |
| 82 | + build_api_dir = os.path.join(build_dir, package, "api") |
| 83 | + build_init_fname = os.path.join(build_dir, package, "__init__.py") |
| 84 | + build_api_init_fname = os.path.join(build_api_dir, "__init__.py") |
| 85 | + try: |
| 86 | + os.chdir(build_dir) |
| 87 | + # Generates `keras_cv/api` directory. |
| 88 | + if os.path.exists(build_api_dir): |
| 89 | + shutil.rmtree(build_api_dir) |
| 90 | + if os.path.exists(build_init_fname): |
| 91 | + os.remove(build_init_fname) |
| 92 | + os.makedirs(build_api_dir) |
| 93 | + namex.generate_api_files( |
| 94 | + "keras_cv", code_directory="src", target_directory="api" |
| 95 | + ) |
| 96 | + # Creates `keras_cv/__init__.py` importing from `keras_cv/api` |
| 97 | + update_package_init(build_init_fname) |
| 98 | + # Add __version__ to keras package |
| 99 | + export_version_string(build_api_init_fname) |
| 100 | + # Copy back the keras_cv/api and keras_cv/__init__.py from build dir |
| 101 | + if os.path.exists(code_api_dir): |
| 102 | + shutil.rmtree(code_api_dir) |
| 103 | + shutil.copytree(build_api_dir, code_api_dir) |
| 104 | + shutil.copy(build_init_fname, code_init_fname) |
| 105 | + finally: |
| 106 | + # Clean up: remove the build directory (no longer needed) |
| 107 | + shutil.rmtree(build_dir) |
| 108 | + |
| 109 | + |
| 110 | +if __name__ == "__main__": |
| 111 | + build() |
0 commit comments