Use temp dir and abs path in api_gen.py

README.md

@@ -50,6 +50,12 @@ pip install -r requirements.txt
 python pip_build.py --install
 ```
 
+3. Run API generation script when creating PRs that update `keras_export` public APIs:
+
+```
+./shell/api_gen.sh
+```
+
 #### Adding GPU support
 
 The `requirements.txt` file will install a CPU-only version of TensorFlow, JAX, and PyTorch. For GPU support, we also

api_gen.py

@@ -18,28 +18,41 @@ def ignore_files(_, filenames):
     return [f for f in filenames if f.endswith("_test.py")]
 
 
-def create_legacy_directory():
-    API_DIR = os.path.join(package, "api")
+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, "tmp_build_dir")
+    if os.path.exists(build_dir):
+        shutil.rmtree(build_dir)
+    os.mkdir(build_dir)
+    shutil.copytree(
+        package, os.path.join(build_dir, package), ignore=ignore_files
+    )
+    return build_dir
+
+
+def create_legacy_directory(package_dir):
+    src_dir = os.path.join(package_dir, "src")
+    api_dir = os.path.join(package_dir, "api")
     # Make keras/_tf_keras/ by copying keras/
-    tf_keras_dirpath_parent = os.path.join(API_DIR, "_tf_keras")
+    tf_keras_dirpath_parent = os.path.join(api_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.api._tf_keras import keras\n")
-    with open(os.path.join(API_DIR, "__init__.py")) as f:
+    with open(os.path.join(api_dir, "__init__.py")) as f:
         init_file = f.read()
         init_file = init_file.replace(
             "from keras.api import _legacy",
             "from keras.api import _tf_keras",
         )
-    with open(os.path.join(API_DIR, "__init__.py"), "w") as f:
+    with open(os.path.join(api_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.api 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(API_DIR):
-        dirpath = os.path.join(API_DIR, dirname)
+    for dirname in os.listdir(api_dir):
+        dirpath = os.path.join(api_dir, dirname)
         if os.path.isdir(dirpath) and dirname not in (
             "_legacy",
             "_tf_keras",
@@ -57,16 +70,16 @@ def create_legacy_directory():
     # Copy keras/_legacy/ file contents to keras/_tf_keras/keras
     legacy_submodules = [
         path[:-3]
-        for path in os.listdir(os.path.join(package, "src", "legacy"))
+        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(package, "src", "legacy"))
-        if os.path.isdir(os.path.join(package, "src", "legacy", 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, "_legacy")):
+    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)
@@ -102,19 +115,18 @@ def create_legacy_directory():
                     f.write(legacy_contents)
 
     # Delete keras/api/_legacy/
-    shutil.rmtree(os.path.join(API_DIR, "_legacy"))
+    shutil.rmtree(os.path.join(api_dir, "_legacy"))
 
 
-def export_version_string():
-    API_INIT = os.path.join(package, "api", "__init__.py")
-    with open(API_INIT) as f:
+def export_version_string(api_init_fname):
+    with open(api_init_fname) as f:
         contents = f.read()
-    with open(API_INIT, "w") as f:
+    with open(api_init_fname, "w") as f:
         contents += "from keras.src.version import __version__\n"
         f.write(contents)
 
 
-def update_package_init():
+def update_package_init(init_fname):
     contents = """
 # Import everything from /api/ into keras.
 from keras.api import *  # noqa: F403
@@ -142,34 +154,48 @@ def __dir__():
     for name in globals().keys()
     if not (name.startswith("_") or name in ("src", "api"))
 ]"""
-    with open(os.path.join(package, "__init__.py")) as f:
+    with open(init_fname) as f:
         init_contents = f.read()
-    with open(os.path.join(package, "__init__.py"), "w") as f:
+    with open(init_fname, "w") as f:
         f.write(init_contents.replace("\nfrom keras import api", contents))
 
 
-if __name__ == "__main__":
+def build():
     # Backup the `keras/__init__.py` and restore it on error in api gen.
-    os.makedirs(os.path.join(package, "api"), exist_ok=True)
-    init_fname = os.path.join(package, "__init__.py")
-    backup_init_fname = os.path.join(package, "__init__.py.bak")
+    root_path = os.path.dirname(os.path.abspath(__file__))
+    code_api_dir = os.path.join(root_path, package, "api")
+    code_init_fname = os.path.join(root_path, package, "__init__.py")
+    # Create temp build dir
+    build_dir = copy_source_to_build_directory(root_path)
+    build_api_dir = os.path.join(build_dir, package, "api")
+    build_init_fname = os.path.join(build_dir, package, "__init__.py")
+    build_api_init_fname = os.path.join(build_api_dir, "__init__.py")
     try:
-        if os.path.exists(init_fname):
-            shutil.move(init_fname, backup_init_fname)
+        os.chdir(build_dir)
         # Generates `keras/api` directory.
+        if os.path.exists(build_api_dir):
+            shutil.rmtree(build_api_dir)
+        if os.path.exists(build_init_fname):
+            os.remove(build_init_fname)
+        os.makedirs(build_api_dir)
         namex.generate_api_files(
             "keras", code_directory="src", target_directory="api"
         )
         # Creates `keras/__init__.py` importing from `keras/api`
-        update_package_init()
-    except Exception as e:
-        if os.path.exists(backup_init_fname):
-            shutil.move(backup_init_fname, init_fname)
-        raise e
+        update_package_init(build_init_fname)
+        # Add __version__ to keras package
+        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(code_api_dir):
+            shutil.rmtree(code_api_dir)
+        shutil.copytree(build_api_dir, code_api_dir)
+        shutil.copy(build_init_fname, code_init_fname)
     finally:
-        if os.path.exists(backup_init_fname):
-            os.remove(backup_init_fname)
-    # Add __version__ to keras package
-    export_version_string()
-    # Creates `_tf_keras` with full keras API
-    create_legacy_directory()
+        # Clean up: remove the build directory (no longer needed)
+        shutil.rmtree(build_dir)
+
+
+if __name__ == "__main__":
+    build()