From 7ee18a1efe963b9c5e532c87ae88fbe01a621f7c Mon Sep 17 00:00:00 2001
From: Ramesh Sampath <1437573+sampathweb@users.noreply.github.com>
Date: Wed, 3 Jan 2024 23:58:31 +0530
Subject: [PATCH] align pip_build with keras (#2256)

* align pip_build with keras

* align pip_build with keras
---
 pip_build.py | 87 ++++++++++++++++++++++++++++++----------------------
 1 file changed, 50 insertions(+), 37 deletions(-)

diff --git a/pip_build.py b/pip_build.py
index 24ad3799dd..29f574001f 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -47,7 +47,7 @@
 
 
 def ignore_files(_, filenames):
-    return [f for f in filenames if "test" in f]
+    return [f for f in filenames if "_test" in f]
 
 
 def export_version_string(version, is_nightly=False):
@@ -71,54 +71,67 @@ def export_version_string(version, is_nightly=False):
         f.write(init_contents + "\n\n" + f'__version__ = "{version}"\n')
 
 
+def copy_source_to_build_directory(root_path):
+    # Copy sources (`keras_cv/` directory and setup files) to build
+    # directory
+    os.chdir(root_path)
+    os.mkdir(build_directory)
+    shutil.copytree(
+        package, os.path.join(build_directory, package), ignore=ignore_files
+    )
+    for fname in to_copy:
+        shutil.copy(fname, os.path.join(f"{build_directory}", fname))
+    os.chdir(build_directory)
+
+
+def run_namex_conversion():
+    # Restructure the codebase so that source files live in `keras_cv/src`
+    namex.convert_codebase(package, code_directory="src")
+
+    # Generate API __init__.py files in `keras_cv/`
+    namex.generate_api_files(package, code_directory="src", verbose=True)
+
+
+def build_and_save_output(root_path, __version__):
+    """Build the package."""
+    os.system("python3 -m build")
+
+    # Save the dist files generated by the build process
+    os.chdir(root_path)
+    if not os.path.exists(dist_directory):
+        os.mkdir(dist_directory)
+    for fpath in glob.glob(
+        os.path.join(build_directory, dist_directory, "*.*")
+    ):
+        shutil.copy(fpath, dist_directory)
+
+    # Find the .whl file path
+    whl_path = None
+    for fname in os.listdir(dist_directory):
+        if __version__ in fname and fname.endswith(".whl"):
+            whl_path = os.path.abspath(os.path.join(dist_directory, fname))
+    if whl_path:
+        print(f"Build successful. Wheel file available at {whl_path}")
+    else:
+        print("Build failed.")
+    return whl_path
+
+
 def build(root_path, is_nightly=False):
     if os.path.exists(build_directory):
         raise ValueError(f"Directory already exists: {build_directory}")
 
-    whl_path = None
     try:
-        # Copy sources (`keras_cv/` directory and setup files) to build dir
-        os.chdir(root_path)
-        os.mkdir(build_directory)
-        shutil.copytree(
-            package, os.path.join(build_directory, package), ignore=ignore_files
-        )
-        for fname in to_copy:
-            shutil.copy(fname, os.path.join(f"{build_directory}", fname))
-        os.chdir(build_directory)
-
-        # Restructure the codebase so that source files live in `keras_cv/src`
-        namex.convert_codebase(package, code_directory="src")
-
-        # Generate API __init__.py files in `keras_cv/`
-        namex.generate_api_files(package, code_directory="src", verbose=True)
-
+        copy_source_to_build_directory(root_path)
+        run_namex_conversion()
         # Make sure to export the __version__ string
         from keras_cv.src import __version__  # noqa: E402
 
         export_version_string(__version__, is_nightly)
-
-        # Build the package
-        os.system("python3 -m build")
-
-        # Save the dist files generated by the build process
-        os.chdir(root_path)
-        if not os.path.exists(dist_directory):
-            os.mkdir(dist_directory)
-        for fpath in glob.glob(
-            os.path.join(build_directory, dist_directory, "*.*")
-        ):
-            shutil.copy(fpath, dist_directory)
-
-        # Find the .whl file path
-        for fname in os.listdir(dist_directory):
-            if __version__ in fname and fname.endswith(".whl"):
-                whl_path = os.path.abspath(os.path.join(dist_directory, fname))
-        print(f"Build successful. Wheel file available at {whl_path}")
+        return build_and_save_output(root_path, __version__)
     finally:
         # Clean up: remove the build directory (no longer needed)
         shutil.rmtree(build_directory)
-    return whl_path
 
 
 def install_whl(whl_fpath):