Add various functions to onboard_project (#846)

* Added various args to onboard_project

-Added zip file support with --zip-password arg
-Added --clean-project arg to run clean_project before extraction
-Added --timestamp arg to append a timestamp to the project name
-Added --wildebeest arg to run a Wildebeest analysis on the extracted project

* Remove --zip-password, switch to password prompt

* Rename variables and methods

Author: TaperChipmunk32

poetry.lock

@@ -100,6 +100,7 @@ openpyxl = "^3.1.2"
 packaging = "^22.0"
 hanzidentifier = "^1.2.0"
 jarowinkler = "^2.0.1"
+wildebeest-nlp = ">=0.9.0"
 
 [tool.poetry.group.dev.dependencies]
 types-pyyaml = "^6.0.12.12"

pyproject.toml

@@ -3,16 +3,19 @@
 import shutil
 from datetime import datetime
 from pathlib import Path
+
 from tqdm import tqdm
 
+
 def parse_arguments():
     parser = argparse.ArgumentParser(description="Clean up unnecessary files and folders.")
-    parser.add_argument("--input", type=Path, default=Path("S:/Paratext/projects"), help="Folder to search.")
+    parser.add_argument("--input", type=Path, default=Path("/root/M/Paratext/projects"), help="Folder to search.")
     parser.add_argument("--delete-subfolders", action="store_true", help="Delete subfolders inside projects.")
     parser.add_argument("--confirm-delete", action="store_true", help="Skip confirmation and delete directly.")
     parser.add_argument("--dry-run", action="store_true", help="Generate a CSV report without deleting.")
     return parser.parse_args()
 
+
 def should_delete(path: Path):
     patterns = [
         "Notes",
@@ -26,6 +29,7 @@ def should_delete(path: Path):
     ]
     return any(pattern in path.name for pattern in patterns)
 
+
 def find_items_to_delete(root_path: Path, delete_subfolders: bool):
     files_to_delete = []
     folders_to_delete = []
@@ -36,19 +40,20 @@ def find_items_to_delete(root_path: Path, delete_subfolders: bool):
             print(f" Warning: Ignoring symlink found: {project_folder}")
             continue
         if project_folder.is_dir():
-            for path in project_folder.glob('*'):
+            for path in project_folder.glob("*"):
                 if path.is_file() and should_delete(path):
                     files_to_delete.append(path)
                 if path.is_dir() and delete_subfolders:
                     folders_to_delete.append(path)
 
     return files_to_delete, folders_to_delete
 
-def execute_and_report(args):
+
+def clean_projects(args):
     now = datetime.now()
-    now_filestamp = now.strftime('%Y%m%d_%H%M%S')
-    now_csv_date = now.strftime('%Y %m %d')
-    now_csv_time = now.strftime('%H:%M:%S')
+    now_filestamp = now.strftime("%Y%m%d_%H%M%S")
+    now_csv_date = now.strftime("%Y %m %d")
+    now_csv_time = now.strftime("%H:%M:%S")
 
     # Find files/folders to delete, with subfolder handling based on the option
     files_to_delete, folders_to_delete = find_items_to_delete(args.input, args.delete_subfolders)
@@ -64,7 +69,7 @@ def execute_and_report(args):
         csv_writer.writerow(["Path", "Type", "Size (bytes)", "Deleted"])
 
         for folder_to_delete in folders_to_delete:
-            size = sum(f.stat().st_size for f in folder_to_delete.glob('*') if f.is_file())
+            size = sum(f.stat().st_size for f in folder_to_delete.glob("*") if f.is_file())
             total_size += size
 
             deleted = "No" if args.dry_run else try_delete(folder_to_delete, args)
@@ -88,23 +93,25 @@ def delete_item(item):
         shutil.rmtree(item)
     return
 
+
 def try_delete(item: Path, args) -> str:
 
     if args.confirm_delete:
         delete_item(item)
         return "Yes"
     else:
         confirmation = input(f"Delete {item}? (y/n): ").strip().lower()
-        if confirmation == 'y':
+        if confirmation == "y":
             delete_item(item)
             return "Yes"
         else:
             return "Skipped"
 
 
 def main():
-    args = parse_arguments()   
-    execute_and_report(args)
+    args = parse_arguments()
+    clean_projects(args)
+
 
 if __name__ == "__main__":
     main()

scripts/clean_projects.py

@@ -1,9 +1,16 @@
 import argparse
+import getpass
 import logging
+import tempfile
+import zipfile
+from datetime import datetime
 from pathlib import Path
 
+import wildebeest.wb_analysis as wb_ana
 import yaml
 
+from scripts.clean_projects import clean_projects
+
 from .collect_verse_counts import collect_verse_counts
 from .environment import SIL_NLP_ENV
 from .extract_corpora import extract_corpora
@@ -58,6 +65,50 @@ def copy_paratext_project_folder(source_dir: Path, project_name: str, overwrite=
             _copy_file_to_paratext_project(source_item, target_item, overwrite=overwrite)
 
 
+def collect_verse_counts_wrapper(project_name: str, verse_counts_config: dict) -> None:
+
+    output_folder = Path(
+        verse_counts_config.get("output_folder", SIL_NLP_ENV.mt_experiments_dir / "verse_counts" / project_name)
+    )
+    if not output_folder.exists():
+        output_folder.mkdir(parents=True, exist_ok=True)
+
+    input_folder = verse_counts_config.get("input_folder", SIL_NLP_ENV.mt_scripture_dir)
+
+    file_patterns = verse_counts_config.get("files", f"*{project_name}*.txt")
+
+    input_folder_path = Path(input_folder)
+    if not input_folder_path.exists():
+        LOGGER.error(f"Input folder '{input_folder_path}' does not exist. Skipping verse counts collection.")
+        return
+
+    matched_files = list(input_folder_path.glob(file_patterns))
+    if not matched_files:
+        LOGGER.error(
+            f"No files matching pattern '{file_patterns}' found in '{input_folder_path}'. Skipping verse counts collection."
+        )
+        return
+
+    collect_verse_counts(
+        input_folder=input_folder_path,
+        output_folder=output_folder,
+        file_patterns=file_patterns,
+        deutero=verse_counts_config.get("deutero", False),
+        recount=verse_counts_config.get("recount", False),
+    )
+
+
+def get_config(config_path: str) -> dict:
+    if config_path:
+        config_file = Path(config_path)
+        if not config_file.exists():
+            raise FileNotFoundError(f"Config file '{config_file}' does not exist.")
+        with config_file.open("r", encoding="utf-8") as file:
+            return yaml.safe_load(file)
+    else:
+        return {}
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Performs several steps to onboard a new project before training a model.",
@@ -97,84 +148,95 @@ def main() -> None:
         action="store_true",
         help="Collect various counts from the extracted Paratext project.",
     )
+    parser.add_argument(
+        "--clean-project",
+        default=False,
+        action="store_true",
+        help="Cleans the Paratext project folder by removing unnecessary files and folders.",
+    )
+    parser.add_argument(
+        "--timestamp",
+        default=False,
+        action="store_true",
+        help="Add a timestamp to the project folder name when creating a new Paratext project folder.",
+    )
+    parser.add_argument(
+        "--wildebeest", default=False, action="store_true", help="Run Wildebeest analysis on the extracted corpora."
+    )
 
     args = parser.parse_args()
     if not args.project:
         raise ValueError("Project name is required. Please provide a valid Paratext project name using <project>.")
 
+    config = get_config(args.config) if args.config else {}
+
+    if args.project.endswith(".zip"):
+        with zipfile.ZipFile(args.project, "r") as zip_ref:
+            # Check if any file in the zip is encrypted
+            temp_dir = tempfile.TemporaryDirectory()
+            needs_password = any(zinfo.flag_bits & 0x1 for zinfo in zip_ref.infolist())
+            if needs_password:
+                pwd = getpass.getpass(prompt=f"Enter password for zip file '{args.project}': ")
+                zip_ref.extractall(temp_dir.name, pwd=pwd.encode())
+            else:
+                zip_ref.extractall(temp_dir.name)
+        args.copy_from = temp_dir.name
+        args.project = Path(args.project).stem
+
     project_name = args.project
+    if args.timestamp:
+
+        now = datetime.now()
+        timestamp = now.strftime("%Y_%m_%d")
+        project_name = f"{args.project}_{timestamp}"
+        LOGGER.info(f"Timestamping project. New project name: {project_name}")
 
     if args.copy_from:
         LOGGER.info(f"Onboarding project: {args.project}")
         paratext_project_dir: Path = create_paratext_project_folder_if_not_exists(project_name)
         copy_paratext_project_folder(Path(args.copy_from), paratext_project_dir, overwrite=args.overwrite)
 
-    if args.config:
-        config_file = Path(args.config)
-        if not config_file.exists():
-            raise FileNotFoundError(f"Config file '{config_file}' does not exist.")
-        with config_file.open("r", encoding="utf-8") as file:
-            config = yaml.safe_load(file)
-    else:
-        raise ValueError("Config file is required. Please provide a valid configuration file using --config.")
+    if args.clean_project:
+        LOGGER.info(f"Cleaning Paratext project folder for {project_name}.")
+        clean_projects(
+            argparse.Namespace(
+                input=get_paratext_project_dir(project_name),
+                delete_subfolders=True,
+                confirm_delete=True,
+                dry_run=False,
+            )
+        )
 
     if args.extract_corpora:
-        LOGGER.info(f"Extracting {project_name}.")
+        extract_config = config.get("extract_corpora", {})
         extract_corpora(
             projects={project_name},
-            books_to_include=config["extract_corpora"]["include"] if "include" in config["extract_corpora"] else [],
-            books_to_exclude=config["extract_corpora"]["exclude"] if "exclude" in config["extract_corpora"] else [],
-            include_markers=(config["extract_corpora"]["markers"] if "markers" in config["extract_corpora"] else False),
-            extract_lemmas=config["extract_corpora"]["lemmas"] if "lemmas" in config["extract_corpora"] else False,
-            extract_project_vrefs=(
-                config["extract_corpora"]["project-vrefs"] if "project-vrefs" in config["extract_corpora"] else False
-            ),
+            books_to_include=extract_config.get("include", []),
+            books_to_exclude=extract_config.get("exclude", []),
+            include_markers=extract_config.get("markers", False),
+            extract_lemmas=extract_config.get("lemmas", False),
+            extract_project_vrefs=extract_config.get("project-vrefs", False),
         )
 
     if args.collect_verse_counts:
         if not args.extract_corpora:
             LOGGER.warning(
                 "--extract_corpora was not included. Collecting verse counts requires the corpus to be extracted first."
             )
-
         LOGGER.info(f"Collecting verse counts from {project_name}.")
+        collect_verse_counts_wrapper(project_name, config.get("verse_counts", {}))
 
-        if config["verse_counts"]["output_folder"]:
-            output_folder = Path(config["verse_counts"]["output_folder"])
-            if not output_folder.exists():
-                output_folder.mkdir(parents=True, exist_ok=True)
-        else:
-            output_folder = SIL_NLP_ENV.mt_experiments_dir / "verse_counts" / project_name
-            if not output_folder.exists():
-                output_folder.mkdir(parents=True, exist_ok=True)
-        input_folder = (
-            config["verse_counts"]["input_folder"]
-            if "input_folder" in config["verse_counts"]
-            else SIL_NLP_ENV.mt_scripture_dir
-        )
-        file_patterns = (
-            config["verse_counts"]["files"] if "files" in config["verse_counts"] else f"*{project_name}*.txt"
-        )
-
-        input_folder_path = Path(input_folder)
-        if not input_folder_path.exists():
-            LOGGER.error(f"Input folder '{input_folder_path}' does not exist. Skipping verse counts collection.")
-            return
-
-        matched_files = list(input_folder_path.glob(file_patterns))
-        if not matched_files:
-            LOGGER.error(
-                f"No files matching pattern '{file_patterns}' found in '{input_folder_path}'. Skipping verse counts collection."
-            )
-            return
-
-        collect_verse_counts(
-            input_folder=input_folder_path,
-            output_folder=output_folder,
-            file_patterns=file_patterns,
-            deutero=config["verse_counts"]["deutero"] if "deutero" in config["verse_counts"] else False,
-            recount=config["verse_counts"]["recount"] if "recount" in config["verse_counts"] else False,
-        )
+    if args.wildebeest:
+        if not args.extract_corpora:
+            LOGGER.warning("--extract_corpora was not included. Wildebeest requires the corpus to be extracted first.")
+
+        extract_file = list(SIL_NLP_ENV.mt_scripture_dir.glob(f"*{project_name}.txt"))[0]
+        LOGGER.info(f"Running Wildebeest analysis on {extract_file}.")
+        with (
+            open(f"{project_name}_wildebeest.json", "w", encoding="utf-8") as json_f,
+            open(f"{project_name}_wildebeest.txt", "w", encoding="utf-8") as txt_f,
+        ):
+            wb_ana.process(in_file=extract_file, json_output=json_f, pp_output=txt_f)
 
 
 if __name__ == "__main__":