refactor: add typing

This commit is contained in:
RedDeadDepresso
2024-09-25 00:59:05 +01:00
parent 8ec39fe30d
commit c797bff858
4 changed files with 50 additions and 38 deletions

View File

@@ -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")

View File

@@ -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()

View File

@@ -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()

View File

@@ -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