Files
keras-team--keras/api_gen.py
T
2026-07-13 12:20:15 +08:00

188 lines
7.4 KiB
Python

"""Script to generate keras public API in `keras/api` directory.
Usage:
Run via `./shell/api_gen.sh`.
It generates API and formats user and generated APIs.
"""
import os
import re
import shutil
import namex
PACKAGE = "keras"
BUILD_DIR_NAME = "tmp_build_dir"
def ignore_files(_, filenames):
return [f for f in filenames if f.endswith("_test.py")]
def copy_source_to_build_directory(root_path):
# Copy sources (`keras/` directory and setup files) to build dir
build_dir = os.path.join(root_path, BUILD_DIR_NAME)
build_package_dir = os.path.join(build_dir, PACKAGE)
build_src_dir = os.path.join(build_package_dir, "src")
root_src_dir = os.path.join(root_path, PACKAGE, "src")
if os.path.exists(build_dir):
shutil.rmtree(build_dir)
os.makedirs(build_package_dir)
shutil.copytree(root_src_dir, build_src_dir)
return build_dir
def create_legacy_directory(package_dir):
src_dir = os.path.join(package_dir, "src")
# Make keras/_tf_keras/ by copying keras/
tf_keras_dirpath_parent = os.path.join(package_dir, "_tf_keras")
tf_keras_dirpath = os.path.join(tf_keras_dirpath_parent, "keras")
os.makedirs(tf_keras_dirpath, exist_ok=True)
with open(os.path.join(tf_keras_dirpath_parent, "__init__.py"), "w") as f:
f.write("from keras._tf_keras import keras\n")
with open(os.path.join(package_dir, "__init__.py")) as f:
init_file = f.read()
init_file = init_file.replace(
"from keras import _legacy as _legacy",
"from keras import _tf_keras as _tf_keras",
)
with open(os.path.join(package_dir, "__init__.py"), "w") as f:
f.write(init_file)
# Remove the import of `_tf_keras` in `keras/_tf_keras/keras/__init__.py`
init_file = init_file.replace("from keras import _tf_keras\n", "\n")
with open(os.path.join(tf_keras_dirpath, "__init__.py"), "w") as f:
f.write(init_file)
for dirname in os.listdir(package_dir):
dirpath = os.path.join(package_dir, dirname)
if os.path.isdir(dirpath) and dirname not in (
"_legacy",
"_tf_keras",
"src",
):
destpath = os.path.join(tf_keras_dirpath, dirname)
if os.path.exists(destpath):
shutil.rmtree(destpath)
shutil.copytree(
dirpath,
destpath,
ignore=ignore_files,
)
# Copy keras/_legacy/ file contents to keras/_tf_keras/keras
legacy_submodules = [
path[:-3]
for path in os.listdir(os.path.join(src_dir, "legacy"))
if path.endswith(".py")
]
legacy_submodules += [
path
for path in os.listdir(os.path.join(src_dir, "legacy"))
if os.path.isdir(os.path.join(src_dir, "legacy", path))
]
for root, _, fnames in os.walk(os.path.join(package_dir, "_legacy")):
for fname in fnames:
if fname.endswith(".py"):
legacy_fpath = os.path.join(root, fname)
tf_keras_root = root.replace(
os.path.join(os.path.sep, "_legacy"),
os.path.join(os.path.sep, "_tf_keras", "keras"),
)
core_api_fpath = os.path.join(
root.replace(os.path.join(os.path.sep, "_legacy"), ""),
fname,
)
if not os.path.exists(tf_keras_root):
os.makedirs(tf_keras_root)
tf_keras_fpath = os.path.join(tf_keras_root, fname)
with open(legacy_fpath) as f:
legacy_contents = f.read()
legacy_contents = legacy_contents.replace(
"keras._legacy", "keras._tf_keras.keras"
)
if os.path.exists(core_api_fpath):
with open(core_api_fpath) as f:
core_api_contents = f.read()
core_api_contents = core_api_contents.replace(
"from keras import _tf_keras as _tf_keras\n", ""
)
for legacy_submodule in legacy_submodules:
core_api_contents = core_api_contents.replace(
f"from keras import {legacy_submodule} as {legacy_submodule}\n", # noqa: E501
"",
)
core_api_contents = core_api_contents.replace(
f"keras.{legacy_submodule}",
f"keras._tf_keras.keras.{legacy_submodule}",
)
# Remove duplicate generated comments string.
legacy_contents = re.sub(r"\n", r"\\n", legacy_contents)
legacy_contents = re.sub('""".*"""', "", legacy_contents)
legacy_contents = re.sub(r"\\n", r"\n", legacy_contents)
# If the same module is in legacy and core_api, use legacy
legacy_imports = re.findall(
r"import (\w+)", legacy_contents
)
for import_name in legacy_imports:
core_api_contents = re.sub(
f"\n.* import {import_name} as {import_name}\n",
r"\n",
core_api_contents,
)
legacy_contents = f"{core_api_contents}\n{legacy_contents}"
with open(tf_keras_fpath, "w") as f:
f.write(legacy_contents)
# Delete keras/api/_legacy/
shutil.rmtree(os.path.join(package_dir, "_legacy"))
def export_version_string(api_init_fname):
with open(api_init_fname) as f:
contents = f.read()
with open(api_init_fname, "w") as f:
contents += "from keras.src.version import __version__ as __version__\n"
f.write(contents)
def build():
root_path = os.path.dirname(os.path.abspath(__file__))
code_api_dir = os.path.join(root_path, PACKAGE, "api")
# Create temp build dir
build_dir = copy_source_to_build_directory(root_path)
build_api_dir = os.path.join(build_dir, PACKAGE)
build_src_dir = os.path.join(build_api_dir, "src")
build_api_init_fname = os.path.join(build_api_dir, "__init__.py")
try:
os.chdir(build_dir)
open(build_api_init_fname, "w").close()
namex.generate_api_files(
"keras",
code_directory="src",
exclude_directories=[
os.path.join("src", "backend", "jax"),
os.path.join("src", "backend", "openvino"),
os.path.join("src", "backend", "tensorflow"),
os.path.join("src", "backend", "torch"),
],
)
# Add __version__ to `api/`.
export_version_string(build_api_init_fname)
# Creates `_tf_keras` with full keras API
create_legacy_directory(package_dir=os.path.join(build_dir, PACKAGE))
# Copy back the keras/api and keras/__init__.py from build directory
if os.path.exists(build_src_dir):
shutil.rmtree(build_src_dir)
if os.path.exists(code_api_dir):
shutil.rmtree(code_api_dir)
shutil.copytree(
build_api_dir, code_api_dir, ignore=shutil.ignore_patterns("src/")
)
finally:
# Clean up: remove the build directory (no longer needed)
shutil.rmtree(build_dir)
if __name__ == "__main__":
build()