diff --git a/modules/fc_kks.py b/modules/fc_kks.py index 6a8f3b7..38b8941 100644 --- a/modules/fc_kks.py +++ b/modules/fc_kks.py @@ -1,8 +1,17 @@ import shutil +from enum import Enum from pathlib import Path +from typing import Literal from util.logger import logger +class CardType(Enum): + UNKNOWN = "UNKNOWN" + KK = "KK" + KKSP = "KKSP" + KKS = "KKS" + + class FilterConvertKKS: def __init__(self, config, file_manager): """Initializes the FilterConvertKKS module. @@ -14,27 +23,27 @@ class FilterConvertKKS: self.file_manager = file_manager self.convert = self.config.fc_kks["Convert"] - def get_list(self, folder_path): + def get_list(self, folder_path: Path) -> list[str]: """Get list of PNG files in the folder.""" folder = Path(folder_path) return [str(file) for file in folder.rglob("*.png")] - def check_png(self, card_path): + def check_png(self, card_path: Path) -> Literal[1, 2, 3]: """Check the PNG file and return its type.""" card_path = Path(card_path) with card_path.open("rb") as card: data = card.read() - card_type = 0 + card_type = CardType.UNKNOWN if b"KoiKatuChara" in data: - card_type = 1 + card_type = CardType.KK if b"KoiKatuCharaSP" in data: - card_type = 2 + card_type = CardType.KKSP elif b"KoiKatuCharaSun" in data: - card_type = 3 - logger.info(f"[{card_type}]", f"{card_path}") + card_type = CardType.KKS + logger.info(f"{card_type.value}", f"{card_path.name}") return card_type - def convert_kk(self, card_name, card_path, destination_path): + def convert_kk(self, card_name: str, card_path: Path, destination_path: Path): """Convert KKS card to KK.""" card_path = Path(card_path) # Convert to Path object with card_path.open(mode="rb") as card: @@ -65,10 +74,13 @@ class FilterConvertKKS: count = len(png_list) if count > 0: - logger.info("SCRIPT", "0: unknown / 1: kk / 2: kksp / 3: kks") + logger.info("SCRIPT", "kk: Koikatsu / KKSP: Koikatsu Special / KKS: Koikatsu Sunshine") + logger.line() + logger.info("FOLDER", str(path)) for png in png_list: - if self.check_png(png) == 3: + if self.check_png(png) == CardType.KKS: kks_card_list.append(png) + logger.line() else: logger.success("SCRIPT", "No PNG files found") return @@ -100,4 +112,4 @@ class FilterConvertKKS: else: logger.success("SCRIPT", f"[{count}] cards moved to [{kks_folder}] folder") else: - logger.success("SCRIPT: No KKS cards found") + logger.success("SCRIPT", "No KKS cards found") diff --git a/modules/install_chara.py b/modules/install_chara.py index 1ed679d..2f0b245 100644 --- a/modules/install_chara.py +++ b/modules/install_chara.py @@ -1,6 +1,6 @@ from pathlib import Path -import codecs from util.logger import logger +from typing import Optional class InstallChara: @@ -15,8 +15,8 @@ class InstallChara: self.game_path = self.config.game_path self.input_path = Path(self.config.install_chara["InputPath"]) # Using Path for input path - def resolve_png(self, image_path): - with codecs.open(image_path[0], "rb") as card: + def resolve_png(self, image_path: Path): + with image_path.open("rb") as card: data = card.read() if b"KoiKatuChara" in data: if b"KoiKatuCharaSP" in data or b"KoiKatuCharaSun" in data: @@ -29,7 +29,7 @@ class InstallChara: else: self.file_manager.copy_and_paste("OVERLAYS", image_path, self.game_path["Overlays"]) - def run(self, folder_path=None): + def run(self, folder_path: Optional[Path] = None): if folder_path is None: folder_path = self.input_path folder_path = Path(folder_path) @@ -40,14 +40,14 @@ class InstallChara: file_list, archive_list = self.file_manager.find_all_files(folder_path) for file in file_list: - file_extension = file[2] - match file_extension: + path, size, extension = file + match extension: case ".zipmod": - self.file_manager.copy_and_paste("MODS", file, self.game_path["mods"]) + self.file_manager.copy_and_paste("MODS", path, self.game_path["mods"]) case ".png": - self.resolve_png(file) + self.resolve_png(path) case _: - basename = Path(file[0]).name + basename = Path(path).name logger.error("UNKNOWN", f"Cannot classify {basename}") logger.line() diff --git a/modules/remove_chara.py b/modules/remove_chara.py index 5cb15a7..d9ff49f 100644 --- a/modules/remove_chara.py +++ b/modules/remove_chara.py @@ -1,4 +1,4 @@ -import codecs +from pathlib import Path from util.logger import logger @@ -14,8 +14,8 @@ class RemoveChara: self.game_path = self.config.game_path self.input_path = self.config.remove_chara["InputPath"] - def resolve_png(self, image_path): - with codecs.open(image_path[0], "rb") as card: + def resolve_png(self, image_path: Path): + with image_path.open("rb") as card: data = card.read() if b"KoiKatuChara" in data: if b"KoiKatuCharaSP" in data or b"KoiKatuCharaSun" in data: @@ -32,13 +32,12 @@ class RemoveChara: file_list, archive_list = self.file_manager.find_all_files(self.input_path) - for file in file_list: - extension = file[2] + for path, size, extension in file_list: match extension: case ".zipmod": - self.file_manager.find_and_remove("MODS", file, self.game_path["mods"]) + self.file_manager.find_and_remove("MODS", path, self.game_path["mods"]) case ".png": - self.resolve_png(file) + self.resolve_png(path) case _: pass logger.line() diff --git a/util/file_manager.py b/util/file_manager.py index 422d260..73143dd 100644 --- a/util/file_manager.py +++ b/util/file_manager.py @@ -6,6 +6,7 @@ import subprocess import time from pathlib import Path from util.logger import logger +from typing import Union, Literal class FileManager: @@ -13,7 +14,7 @@ class FileManager: def __init__(self, config): self.config = config - def find_all_files(self, directory): + def find_all_files(self, directory: Union[Path, str]) -> tuple[Path, int, str]: """Find all files in the given directory.""" directory = Path(directory) file_list = [] @@ -33,9 +34,9 @@ class FileManager: archive_list.sort(key=lambda x: x[1]) return file_list, archive_list - def copy_and_paste(self, type, source_path, destination_folder): + def copy_and_paste(self, type: str, source_path: Union[Path, str], destination_folder: Union[str, Path]): """Copy file from source to destination, handling file conflicts.""" - source_path = Path(source_path[0]) + source_path = Path(source_path) destination_folder = Path(destination_folder) base_name = source_path.name @@ -76,22 +77,22 @@ class FileManager: except Exception as e: logger.error(type, f"An error occurred: {e}") - def find_and_remove(self, type, source_path, destination_folder_path): + def find_and_remove(self, file_type: str, source_path: Union[str, Path], destination_folder: Union[str, Path]): """Remove file if it exists at the destination.""" - source_path = Path(source_path[0]) - destination_folder_path = Path(destination_folder_path) + source_path = Path(source_path) + destination_folder = Path(destination_folder) base_name = source_path.name - destination_path = destination_folder_path / base_name + destination_path = destination_folder / base_name if destination_path.exists(): try: destination_path.unlink() - logger.removed(type, base_name) + logger.removed(file_type, base_name) except OSError as e: - logger.error(type, base_name) + logger.error(file_type, base_name) - def create_archive(self, folders, archive_path): + def create_archive(self, folders: list[Literal["mods", "UserData", "BepInEx"]], archive_path: Union[str, Path]): """Create an archive of the given folders using 7zip.""" # Specify the full path to the 7zip executable path_to_7zip = patoolib.util.find_program("7z") @@ -138,7 +139,7 @@ class FileManager: if process.returncode not in [0, 1]: raise Exception() - def extract_archive(self, archive_path): + def extract_archive(self, archive_path: Union[Path, str]): """Extract the archive.""" archive_path = Path(archive_path) archive_name = archive_path.name