wit module now can simulate FST if the game is a FST. Moved many function from game to wit, added exists method to szs

source/mkw/Game.py

@@ -1,73 +1,29 @@
-import enum
 from pathlib import Path
 
-from source.wt.wit import WITPath
+from source.wt.wit import WITPath, Region
-
-
-class Extension(enum.Enum):
-    """
-    Enum for game extension
-    """
-    FST = ".dol"
-    WBFS = ".wbfs"
-    ISO = ".iso"
-
-    @classmethod
-    def _missing_(cls, value: str) -> "Extension | None":
-        """
-        if not found, search for the same value with lower case
-        :param value: value to search for
-        :return: None if nothing found, otherwise the found value
-        """
-        value = value.lower()
-        for member in filter(lambda m: m.value == value, cls): return member
-        return None
-
-
-class Region(enum.Enum):
-    """
-    Enum for game region
-    """
-    PAL = "PAL"
-    USA = "USA"
-    EUR = "EUR"
-    KOR = "KOR"
 
 
 class Game:
     def __init__(self, path: Path | str):
-        self.path = Path(path) if isinstance(path, str) else path
+        self.wit_path = WITPath(path)
-
-    @property
-    def extension(self) -> Extension:
-        """
-        Returns the extension of the game
-        :return: the extension of the game
-        """
-        return Extension(self.path.suffix)
-
-    @property
-    def id(self) -> str:
-        """
-        Return the id of the game (RMCP01, RMCK01, ...)
-        :return: the id of the game
-        """
-        return WITPath(self.path).analyze()["id6"]
-
-    @property
-    def region(self) -> Region:
-        """
-        Return the region of the game (PAL, USA, EUR, ...)
-        :return: the region of the game
-        """
-        return Region(WITPath(self.path).analyze()["dol_region"])
 
     def is_mkw(self) -> bool:
         """
-        Return True if the game is Mario Kart Wii, else otherwise
+        Return True if the game is Mario Kart Wii, False otherwise
         :return: is the game a MKW game
         """
-        return WITPath(self.path).analyze()["dol_is_mkw"] == 1
+        return self.wit_path.analyze()["dol_is_mkw"] == 1
 
     def is_vanilla(self) -> bool:
-        ...
+        """
+        Return True if the game is vanilla, False if the game is modded
+        :return: if the game is not modded
+        """
+        return not any(self.wit_path[f"./files/rel/lecode-{region.value}.bin"].exists() for region in Region)
+
+    def extract(self, dest: Path | str) -> Path:
+        """
+        Extract the game to the destination directory. If the game is a FST, just copy to the destination
+        :param dest: destination directory
+        """
+        return self.wit_path.extract_all(dest)

source/wt/szs.py

@@ -5,14 +5,18 @@ tools_path = tools_szs_dir / ("wszst.exe" if system == "win64" else "wszst")
 
 
 class SZSPath:
-    __slots__ = ("path",)
+    __slots__ = ("path", "_analyze")
 
     def __init__(self, path: Path | str):
         self.path: Path = path if isinstance(path, Path) else Path(path)
+        self._analyze = None
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<SZSPath: {self.path}>"
 
+    def __eq__(self, other: "SZSPath") -> bool:
+        return self.path == other.path
+
     @better_error(tools_path)
     def _run(self, *args) -> bytes:
         """
@@ -62,10 +66,11 @@ class SZSPath:
 
     def analyze(self) -> dict:
         """
-        Return the analyze of the szs
+        Return the analyze of the file
         :return: dictionnary of key and value of the analyze
         """
-        return self._run_dict("ANALYZE", self.path)
+        if self._analyze is None: self._analyze = self._run_dict("ANALYZE", self.path)
+        return self._analyze
 
     def list_raw(self) -> list[str]:
         """
@@ -75,7 +80,11 @@ class SZSPath:
 
         # cycle though all of the output line of the command, check if the line are empty, and if not,
         # add it to the list. Finally, remove the first line because this is a description of the command
-        return [subfile.strip() for subfile in self._run("list", self.path).decode().splitlines() if subfile][1:]
+        return [
+            subfile.strip()
+            for subfile in self._run("list", self.path).decode().splitlines()
+            if subfile.startswith("./")
+        ]
 
     def list(self) -> list["SZSSubPath"]:
         """
@@ -105,9 +114,12 @@ class SZSSubPath:
         self.szs_path = szs_path
         self.subfile = subfile
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<SZSSubPath: {self.szs_path.path}/{self.subfile}>"
 
+    def __eq__(self, other: "SZSSubPath") -> bool:
+        return self.subfile == other.subfile and self.szs_path == other.szs_path
+
     def extract(self, dest: Path | str) -> Path:
         """
         Extract the subfile to a destination
@@ -124,11 +136,30 @@ class SZSSubPath:
 
         return dest
 
-    def is_dir(self):
+    def exists(self):
+        """
+        Return if the subfile exist in the szs
+        :return: True if the subfile exist, else otherwise
+        """
+        return self in self.szs_path.list()
+
+    def is_dir(self) -> bool:
+        """
+        Return if the subfile is a directory
+        :return: True if the subfile is a directory, else otherwise
+        """
         return self.subfile.endswith("/")
 
-    def is_file(self):
+    def is_file(self) -> bool:
+        """
+        Return if the subfile is a file
+        :return: True if the subfile is a file, else otherwise
+        """
         return not self.is_dir()
 
-    def basename(self):
+    def basename(self) -> str:
+        """
+        Return the basename of the subfile
+        :return: the basename of the subfile
+        """
         return self.subfile.rsplit("/", 1)[-1]

source/wt/wit.py

@@ -1,14 +1,54 @@
+import enum
+import shutil
+
 from source.wt import *
 from source.wt import _run, _run_dict
 
 tools_path = tools_wit_dir / ("wit.exe" if system == "win64" else "wit")
 
 
-class WITPath:
+class Extension(enum.Enum):
-    __slots__ = ("path",)
+    """
+    Enum for game extension
+    """
+    FST = ".dol"
+    WBFS = ".wbfs"
+    ISO = ".iso"
 
-    def __init__(self, path: Path):
+    @classmethod
-        self.path = path
+    def _missing_(cls, value: str) -> "Extension | None":
+        """
+        if not found, search for the same value with lower case
+        :param value: value to search for
+        :return: None if nothing found, otherwise the found value
+        """
+        value = value.lower()
+        for member in filter(lambda m: m.value == value, cls): return member
+        return None
+
+
+class Region(enum.Enum):
+    """
+    Enum for game region
+    """
+    PAL = "PAL"
+    USA = "USA"
+    EUR = "EUR"
+    KOR = "KOR"
+
+
+class WITPath:
+    __slots__ = ("path", "_analyze")
+
+    def __init__(self, path: Path | str):
+        self.path: Path = path if isinstance(path, Path) else Path(path)
+        self._analyze = None
+
+    def __repr__(self) -> str:
+        return f"<WITPath: {self.path}>"
+
+    def __eq__(self, other: "WITPath") -> bool:
+        return self.path == other.path
 
     @better_error(tools_path)
     def _run(self, *args) -> bytes:
@@ -28,9 +68,172 @@ class WITPath:
         """
         return _run_dict(tools_path, *args)
 
+    def _get_fst_root(self) -> Path:
+        """
+        If the game is a FST, return the root of the FST
+        :return: root of the FST
+        """
+        # main.dol is located in ./sys/main.dol, so return parent of parent
+        if self.extension == Extension.FST: return self.path.parent.parent
+
     def analyze(self) -> dict:
         """
         Return the analyze of the file
         :return: dictionnary of key and value of the analyze
         """
-        return self._run_dict("ANALYZE", self.path)
+        if self._analyze is None: self._analyze = self._run_dict("ANALYZE", self.path)
+        return self._analyze
+    
+    def list_raw(self) -> list[str]:
+        """
+        Return the list of subfiles
+        :return: the list of subfiles
+        """
+        if self.extension == Extension.FST:
+            return [
+                str(file.relative_to(self._get_fst_root()))
+                for file in self._get_fst_root().rglob("*")
+            ]
+
+        return [
+            subfile.strip() for subfile
+            in self._run("files", self.path).decode().splitlines()
+            if subfile.startswith("./")
+        ]
+
+    def list(self) -> list["WITSubPath"]:
+        """
+        Return the list of subfiles
+        :return: the list of subfiles
+        """
+        return [self.get_subfile(subfile) for subfile in self.list_raw()]
+
+    def get_subfile(self, subfile: str) -> "WITSubPath":
+        """
+        Return the subfile of the game
+        :return: the subfile
+        """
+        return WITSubPath(self, subfile)
+
+    def __getitem__(self, item):
+        return self.get_subfile(item)
+
+    def __iter__(self):
+        return iter(self.list())
+
+    def extract_all(self, dest: Path | str) -> Path:
+        """
+        Extract all the subfiles to the destination directory
+        :param dest: destination directory
+        :return: the extracted file path
+        """
+        return self["./"].extract(dest, flat=False)
+
+    @property
+    def extension(self) -> Extension:
+        """
+        Returns the extension of the game
+        :return: the extension of the game
+        """
+        return Extension(self.path.suffix)
+
+    @property
+    def id(self) -> str:
+        """
+        Return the id of the game (RMCP01, RMCK01, ...)
+        :return: the id of the game
+        """
+        return self.analyze()["id6"]
+
+    @property
+    def region(self) -> Region:
+        """
+        Return the region of the game (PAL, USA, EUR, ...)
+        :return: the region of the game
+        """
+        return Region(self.analyze()["dol_region"])
+
+
+class WITSubPath:
+    __slots__ = ("wit_path", "subfile")
+
+    def __init__(self, wit_path: WITPath, subfile: str):
+        self.wit_path = wit_path
+        self.subfile = subfile.removeprefix("./").replace("\\", "/")
+
+    def __repr__(self):
+        if self.wit_path.extension == Extension.FST: return f"<WITSubPath: {self._get_fst_path()}>"
+        return f"<WITSubPath: {self.wit_path.path}/{self.subfile}>"
+
+    def __eq__(self, other: "WITSubPath") -> bool:
+        return self.subfile == other.subfile and self.wit_path == other.wit_path
+
+    def _get_fst_path(self) -> Path:
+        """
+        Return the path of the subfile in the FST
+        :return: the path of the subfile in the FST
+        """
+        return self.wit_path._get_fst_root() / self.subfile
+
+    def extract(self, dest: Path | str, flat: bool = True) -> Path:
+        """
+        Extract the subfile to the destination directory
+        :param flat: all files will be extracted directly in the directory, instead of creating subdirectory
+        :param dest: destination directory
+        :return: the extracted file path
+        """
+        dest: Path = dest if isinstance(dest, Path) else Path(dest)
+
+        if self.wit_path.extension == Extension.FST:
+            # if flat is used, extract the file / dir into the destination directory, without subdirectory
+            if flat:
+                os.makedirs(dest, exist_ok=True)
+                # if we are extracting a directory, we need to extract every file recursively
+                if self.is_dir():
+                    for file in (self._get_fst_path()).rglob("*"):
+                        if file.is_file(): shutil.copy(file, dest / file.name)
+                # else we just copy the file
+                else:
+                    shutil.copy(self._get_fst_path(), dest)
+            # if flat is not used, copy the structure of the directory, or just copy the file
+            else:
+                func = shutil.copytree if self.is_dir() else shutil.copy
+                func(self._get_fst_path(), dest / self.subfile)
+
+            return dest / self.basename()
+
+        else:
+            args = []
+            if flat: args.append("--flat")
+            self.wit_path._run("EXTRACT", self.wit_path.path, f"--files=+{self.subfile}", "-d", dest, *args)
+            return dest / self.basename()
+
+    def is_dir(self) -> bool:
+        """
+        Return if the subfile is a directory
+        :return: True if the subfile is a directory, else otherwise
+        """
+        if self.wit_path.extension == Extension.FST:
+            return self._get_fst_path().is_dir()
+        return self.subfile.endswith("/")
+
+    def is_file(self) -> bool:
+        """
+        Return if the subfile is a file
+        :return: True if the subfile is a file, else otherwise
+        """
+        return not self.is_dir()
+
+    def exists(self):
+        """
+        Return if the subfile exist in the game
+        :return: True if the subfile exist, else otherwise
+        """
+        return self in self.wit_path.list()
+
+    def basename(self) -> str:
+        """
+        Return the basename of the subfile
+        :return: the basename of the subfile
+        """
+        return self.subfile.rsplit("/", 1)[-1]