添加清理 pycache 和编译扩展的功能

scripts/compile.py

@@ -6,6 +6,83 @@ from setuptools import setup, Extension
 from Cython.Build import cythonize
 
 
+def delete_source_files(target_dirs):
+    """
+    Deletes the original .py files in the target directories after successful compilation.
+    Caution: This will remove your source code!
+    """
+    if isinstance(target_dirs, str):
+        target_dirs = [target_dirs]
+
+    project_root = os.getcwd()
+    print(f"Deleting source .py files in: {target_dirs}")
+
+    for target_dir in target_dirs:
+        if not os.path.exists(target_dir):
+            continue
+
+        abs_target_path = os.path.abspath(target_dir)
+
+        if os.path.isfile(abs_target_path):
+            if abs_target_path.endswith(".py") and os.path.abspath(
+                abs_target_path
+            ) != os.path.abspath(__file__):
+                print(f"Deleting source: {abs_target_path}")
+                os.remove(abs_target_path)
+            continue
+
+        # If it's a directory, walk through it
+        for root, dirs, files in os.walk(abs_target_path):
+            for file in files:
+                if file.endswith(".py"):
+                    file_path = os.path.join(root, file)
+                    # Skip this script itself
+                    if os.path.abspath(file_path) == os.path.abspath(__file__):
+                        continue
+                    # Skip __init__.py if you want to keep package structure,
+                    # but usually Cython can handle them too.
+                    print(f"Deleting source: {file_path}")
+                    os.remove(file_path)
+
+
+def clean_extensions(target_dirs):
+    """
+    Deletes all compiled .so/.pyd files in the target directories.
+    """
+    if isinstance(target_dirs, str):
+        target_dirs = [target_dirs]
+
+    project_root = os.getcwd()
+    print(f"Cleaning compiled files in: {target_dirs}")
+
+    for target_dir in target_dirs:
+        if not os.path.exists(target_dir):
+            continue
+
+        abs_target_path = os.path.abspath(target_dir)
+
+        if os.path.isfile(abs_target_path):
+            # If it's a .py file, look for its .so/.pyd counterpart
+            dir_name = os.path.dirname(abs_target_path)
+            base_name = os.path.splitext(os.path.basename(abs_target_path))[0]
+            for file in os.listdir(dir_name):
+                if file.startswith(base_name) and (
+                    file.endswith(".so") or file.endswith(".pyd")
+                ):
+                    file_path = os.path.join(dir_name, file)
+                    print(f"Removing: {file_path}")
+                    os.remove(file_path)
+            continue
+
+        # If it's a directory, walk through it
+        for root, dirs, files in os.walk(abs_target_path):
+            for file in files:
+                if file.endswith(".so") or file.endswith(".pyd"):
+                    file_path = os.path.join(root, file)
+                    print(f"Removing: {file_path}")
+                    os.remove(file_path)
+
+
 def build_extensions(target_dirs):
     """
     Compiles all .py files in the target directories (recursively) into .so/.pyd extensions.
@@ -119,17 +196,35 @@ if __name__ == "__main__":
 
     # Check for help flag
     if len(sys.argv) > 1 and sys.argv[1] in ["--help", "-h"]:
-        print("Usage: python scripts/build_extensions.py [directory1] [directory2] ...")
+        print(
+            "Usage: python scripts/build_extensions.py [--clean|--delete-source] [directory1] [directory2] ..."
+        )
         print(
             "Compiles all Python files in the given directories into C extensions (.so/.pyd)."
         )
+        print("Use --clean to remove existing compiled extensions.")
+        print("Use --delete-source to remove source .py files after compilation.")
         print(f"Default directories if none provided: {DEFAULT_TARGETS}")
         sys.exit(0)
 
+    clean_mode = False
+    delete_source_mode = False
     if len(sys.argv) > 1:
-        target_directories = sys.argv[1:]
+        if sys.argv[1] == "--clean":
+            clean_mode = True
+            target_directories = sys.argv[2:]
+        elif sys.argv[1] == "--delete-source":
+            delete_source_mode = True
+            target_directories = sys.argv[2:]
+        else:
+            target_directories = sys.argv[1:]
     else:
         print(f"No directories specified. Using defaults: {DEFAULT_TARGETS}")
         target_directories = DEFAULT_TARGETS
 
-    build_extensions(target_directories)
+    if clean_mode:
+        clean_extensions(target_directories)
+    elif delete_source_mode:
+        delete_source_files(target_directories)
+    else:
+        build_extensions(target_directories)