diff --git a/main.py b/main.py index b92f18c..502bce5 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ import time import os import warnings +import shutil import re from version import __version__ import questionary @@ -19,7 +20,7 @@ parser = argparse.ArgumentParser(prog="sumzero", formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("-V", "--version", action="version", version=f"%(prog)s {__version__}") -parser.add_argument("-c", "--chapter", help="Summarize specific chapters (opens up the chapter selector)", action="store_true") +parser.add_argument("-m", "--merge", help="Merge all of the outputs into a single file", action="store_true") parser.add_argument("-i", "--input", type=str, help="The path to the .txt file of the arc. Required", required=True) parser.add_argument("-o", "--output", type=str, help="Output folder path", default="output") parser.add_argument("-v", "--verbose", help="Verbose mode", action="store_true") @@ -41,6 +42,9 @@ exit() outputdir = os.path.abspath(args.output) if args.output else os.path.join(os.getcwd(), "output") +if args.merge: + originaloutputdir = outputdir + outputdir = os.path.join(outputdir, "temp") os.makedirs(outputdir, exist_ok=True) file = open(args.input, "r", encoding="utf-8") @@ -192,65 +196,61 @@ def handleIndividualChapter(chapter): file.truncate() file.close() +chaptersinarc = [] -# If the user specified a chapter, handle that chapter. Otherwise, summarize all chapters. -if args.chapter: - chaptersinarc = [] +# Get all chapters in the arc and add them to a list +for i in range(len(texts)): + firstline = texts[i].split("\n")[0] - # Get all chapters in the arc and add them to a list - for i in range(len(texts)): - firstline = texts[i].split("\n")[0] + if "Chapter" in firstline and not "Part" in firstline: + chaptersinarc.append(firstline) + +# Ask the user which chapters they want to summarize +chapters: list | None = questionary.checkbox( + "Which chapter(s) do you want to summarize?", + choices=chaptersinarc, +).ask() + +if chapters == None: + print("No chapters selected. Exiting...") + exit() + +# Get the chapter number from the chapter title +for i in range(len(chapters)): + chapters[i] = str(re.search(r"(?<=Chapter )\w+", chapters[i]).group(0)) + +# Sort the chapters in ascending order +chapters.sort() + +print(Fore.YELLOW + "[-] " + "Handling chapter(s) " + ", ".join(chapters)) + +if args.merge: + # Remove the temp folder if it already exists + if os.path.exists(outputdir): + shutil.rmtree(outputdir) + os.mkdir(outputdir) - if "Chapter" in firstline and not "Part" in firstline: - chaptersinarc.append(firstline) - - # Ask the user which chapters they want to summarize - chapters: list | None = questionary.checkbox( - "Which chapter(s) do you want to summarize?", - choices=chaptersinarc, - ).ask() - - if chapters == None: - print("No chapters selected. Exiting...") - exit() - - # Get the chapter number from the chapter title - for i in range(len(chapters)): - chapters[i] = str(re.search(r"(?<=Chapter )\w+", chapters[i]).group(0)) - - # Sort the chapters in ascending order - chapters.sort() - - print(Fore.YELLOW + "[-] " + "Handling chapter(s) " + ", ".join(chapters)) - - # Handle each chapter - for chapter in chapters: - print(Fore.YELLOW + "\n[-] " + "Processing chapter " + chapter + "...") - handleIndividualChapter(chapter.strip()) - print(Fore.GREEN + "[✓] " + "Processed chapter " + chapter + "!") + +# Handle each chapter +for chapter in chapters: + print(Fore.YELLOW + "\n[-] " + "Processing chapter " + chapter + "...") + handleIndividualChapter(chapter.strip()) + print(Fore.GREEN + "[✓] " + "Processed chapter " + chapter + "!") - print(Fore.GREEN + "[✓] " + "Done!") -else: - if args.force and os.path.exists(os.path.join(outputdir, "Summary.txt")): - os.remove(os.path.join(outputdir, "Summary.txt")) - - for i in tqdm.trange(len(texts)): - summary = summarize(i) if not args.dry_run else "Dry run" +# Merge all of the files into a single file +if args.merge: + print(Fore.YELLOW + "\n[-] " + "Merging files...") + with open(os.path.join(originaloutputdir, "Arc 7 Summary.txt"), "w", encoding="utf-8") as outfile: + for chapter in chapters: + with open(os.path.join(outputdir, f"Chapter {chapter} Summary.txt"), "r", encoding="utf-8") as infile: + outfile.write(infile.read()) + outfile.write("\n\n\n") + infile.close() + outfile.close() - if args.verbose: - print('\n------------------------------------------------\nSummary:') - print(summary) - print("------------------------------------------------") - - with open(os.path.join(outputdir, "Summary.txt"), "a", encoding="utf-8") as file: - firstline = texts[i].split("\n")[0] - file.write(f"{firstline}\n{summary}\n\n\n") - - # Remove whitespace - with open(os.path.join(outputdir, "Summary.txt"), "r+", encoding="utf-8") as file: - fart = file.read() - fart = fart.strip() - file.seek(0) - file.write(fart) - file.truncate() - file.close() \ No newline at end of file + # Delete the temp folder + shutil.rmtree(outputdir) + + print(Fore.GREEN + "[✓] " + "Merged files!") + +print(Fore.GREEN + "[✓] " + "Done!") \ No newline at end of file