fix(helpers): type hint, import, typo

Author: isHarryh

UnityPy/helpers/ArchiveStorageManager.py

@@ -1,6 +1,6 @@
 # based on: https://github.com/Razmoth/PGRStudio/blob/master/AssetStudio/PGR/PGR.cs
 import re
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 from ..streams import EndianBinaryReader
 
@@ -10,7 +10,7 @@
     UnityPyBoost = None
 
 UNITY3D_SIGNATURE = b"#$unity3dchina!@"
-DECRYPT_KEY: bytes = None
+DECRYPT_KEY: Optional[bytes] = None
 
 
 def set_assetbundle_decrypt_key(key: Union[bytes, str]):
@@ -66,7 +66,7 @@ class ArchiveStorageDecryptor:
     index: bytes
     substitute: bytes = bytes(0x10)
 
-    def __init__(self, reader: EndianBinaryReader) -> None:
+    def __init__(self, reader: EndianBinaryReader):
         self.unknown_1 = reader.read_u_int()
 
         # read vector data/key vectors
@@ -100,7 +100,6 @@ def __init__(self, reader: EndianBinaryReader) -> None:
         )
 
     def decrypt_block(self, data: bytes, index: int):
-
         if UnityPyBoost:
             return UnityPyBoost.decrypt_block(self.index, self.substitute, data, index)
 
@@ -113,7 +112,7 @@ def decrypt_block(self, data: bytes, index: int):
             index += 1
         return data
 
-    def decrypt_byte(self, view: bytearray, offset: int, index: int):
+    def decrypt_byte(self, view: Union[bytearray, memoryview], offset: int, index: int):
         b = (
             self.substitute[((index >> 2) & 3) + 4]
             + self.substitute[index & 3]
@@ -127,7 +126,7 @@ def decrypt_byte(self, view: bytearray, offset: int, index: int):
         b = view[offset]
         return b, offset + 1, index + 1
 
-    def decrypt(self, data: bytearray, index: int, remaining: int):
+    def decrypt(self, data: Union[bytearray, memoryview], index: int, remaining: int):
         offset = 0
 
         curByte, offset, index = self.decrypt_byte(data, offset, index)

UnityPy/helpers/CompressionHelper.py

@@ -1,11 +1,10 @@
+import brotli
 import gzip
 import lzma
+import lz4.block
 import struct
 from typing import Tuple
 
-import brotli
-import lz4.block
-
 GZIP_MAGIC: bytes = b"\x1f\x8b"
 BROTLI_MAGIC: bytes = b"brotli"
 

UnityPy/helpers/ImportHelper.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
+
 import os
-from typing import Union, List
+from typing import Union, List, Optional, Tuple
 from .CompressionHelper import BROTLI_MAGIC, GZIP_MAGIC
 from ..enums import FileType
 from ..streams import EndianBinaryReader
@@ -18,7 +19,7 @@ def list_all_files(directory: str) -> List[str]:
         val
         for sublist in [
             [os.path.join(dir_path, filename) for filename in filenames]
-            for (dir_path, dirn_ames, filenames) in os.walk(directory)
+            for (dir_path, dirnames, filenames) in os.walk(directory)
             if ".git" not in dir_path
         ]
         for val in sublist
@@ -34,14 +35,14 @@ def find_all_files(directory: str, search_str: str) -> List[str]:
                 for filename in filenames
                 if search_str in filename
             ]
-            for (dir_path, dirn_ames, filenames) in os.walk(directory)
+            for (dir_path, dirnames, filenames) in os.walk(directory)
             if ".git" not in dir_path
         ]
         for val in sublist
     ]
 
 
-def check_file_type(input_) -> Union[FileType, EndianBinaryReader]:
+def check_file_type(input_) -> Tuple[Optional[FileType], Optional[EndianBinaryReader]]:
     if isinstance(input_, str) and os.path.isfile(input_):
         reader = EndianBinaryReader(open(input_, "rb"))
     elif isinstance(input_, EndianBinaryReader):
@@ -124,10 +125,10 @@ def check_file_type(input_) -> Union[FileType, EndianBinaryReader]:
 
 def parse_file(
     reader: EndianBinaryReader,
-    parent,
+    parent: files.File,
     name: str,
-    typ: FileType = None,
-    is_dependency=False,
+    typ: Optional[FileType] = None,
+    is_dependency: bool = False
 ) -> Union[files.File, EndianBinaryReader]:
     if typ is None:
         typ, _ = check_file_type(reader)

UnityPy/helpers/MeshHelper.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
+
 import math
 import struct
-from typing import Optional, List, Tuple, Union, TypeVar
+from typing import Optional, List, Sequence, Tuple, Union, TypeVar
 
-from ..enums.MeshTopology import MeshTopology
 from ..classes.generated import (
     ChannelInfo,
     StreamInfo,
@@ -13,8 +13,7 @@
     Vector3f,
     Vector4f,
 )
-from .PackedBitVector import unpack_floats, unpack_ints
-
+from ..enums.MeshTopology import MeshTopology
 from ..enums.VertexFormat import (
     VertexChannelFormat,
     VertexFormat2017,
@@ -23,6 +22,7 @@
     VERTEX_FORMAT_2017_STRUCT_TYPE_MAP,
     VERTEX_FORMAT_STRUCT_TYPE_MAP,
 )
+from .PackedBitVector import unpack_floats, unpack_ints
 from .ResourceReader import get_resource_data
 
 try:
@@ -37,11 +37,11 @@
 T = TypeVar("T")
 
 
-def flat_list_to_tuples(data: List[T], item_size: int) -> List[tuple[T]]:
+def flat_list_to_tuples(data: Sequence[T], item_size: int) -> List[tuple[T, ...]]:
     return [tuple(data[i : i + item_size]) for i in range(0, len(data), item_size)]
 
 
-def vector_list_to_tuples(data: List[Vector2f, Vector3f, Vector4f]) -> List[tuple]:
+def vector_list_to_tuples(data: List[Union[Vector2f, Vector3f, Vector4f]]) -> List[tuple]:
     if isinstance(data[0], Vector2f):
         return [(v.x, v.y) for v in data]
     elif isinstance(data[0], Vector3f):
@@ -52,7 +52,7 @@ def vector_list_to_tuples(data: List[Vector2f, Vector3f, Vector4f]) -> List[tupl
         raise ValueError("Unknown vector type")
 
 
-def zeros(shape: Tuple[int, ...]) -> list:
+def zeros(shape: Union[Tuple[int], Tuple[int, int]]) -> Union[List, List[List]]:
     if len(shape) == 1:
         return [0] * shape[0]
     elif len(shape) == 2:
@@ -97,7 +97,7 @@ def __init__(
         src: Union[Mesh, SpriteRenderData],
         version: Optional[Tuple[int, int, int, int]] = None,
         endianess: str = "<",
-    ) -> None:
+    ):
         self.src = src
         self.endianess = endianess
         if version is not None:

UnityPy/helpers/PackedBitVector.py

@@ -1,6 +1,5 @@
 from typing import Optional, List, TYPE_CHECKING, Tuple
 
-# from ..objects.math import Quaternionf
 if TYPE_CHECKING:
     from ..classes.generated import PackedBitVector
 

UnityPy/helpers/Tpk.py

@@ -4,7 +4,7 @@
 from importlib.resources import open_binary
 from io import BytesIO
 from struct import Struct
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from .TypeTreeHelper import TypeTreeNode
 
@@ -19,7 +19,9 @@ def init():
 
     global TPKTYPETREE
     with BytesIO(data) as stream:
-        TPKTYPETREE = TpkFile(stream).GetDataBlob()
+        blob = TpkFile(stream).GetDataBlob()
+        assert isinstance(blob, TpkTypeTreeBlob)
+        TPKTYPETREE = blob
 
 
 def get_typetree_node(class_id: int, version: tuple):
@@ -338,8 +340,8 @@ class TpkUnityClass:
     Name: int
     Base: int
     Flags: TpkUnityClassFlags
-    EditorRootNode: int
-    ReleaseRootNode: int
+    EditorRootNode: Optional[int]
+    ReleaseRootNode: Optional[int]
 
     def __init__(self, stream: BytesIO) -> None:
         self.Name, self.Base, Flags = TpkUnityClass.Struct.unpack(
@@ -473,7 +475,7 @@ def Count(self) -> int:
 class TpkCommonString:
     __slots__ = ("VersionInformation", "StringBufferIndices")
     VersionInformation: List[Tuple[UnityVersion, int]]
-    StringBufferIndices: List[int]
+    StringBufferIndices: Tuple[int]
 
     def __init__(self, stream: BytesIO) -> None:
         (versionCount,) = INT32.unpack(stream.read(INT32.size))

UnityPy/helpers/TypeTreeHelper.py

@@ -1,4 +1,4 @@
-from __future__ import annotations
+from __future__ import annotations
 
 import re
 from typing import TYPE_CHECKING, Any, Optional, Union
@@ -49,6 +49,7 @@
     "string": EndianBinaryReader.read_aligned_string,
     "TypelessData": EndianBinaryReader.read_byte_array,
 }
+
 FUNCTION_READ_MAP_ARRAY = {
     "SInt8": EndianBinaryReader.read_byte_array,
     "UInt8": EndianBinaryReader.read_u_byte_array,
@@ -83,7 +84,7 @@ def copy(self) -> TypeTreeConfig:
         return TypeTreeConfig(self.as_dict, self.assetsfile, self.has_registry)
 
 
-def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> TypeTreeNode:
+def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> Optional[TypeTreeNode]:
     typ = ref_object["type"]
     if isinstance(typ, dict):
         cls = typ["class"]

UnityPy/helpers/TypeTreeNode.py

@@ -2,7 +2,7 @@
 
 import re
 from struct import Struct
-from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Any, Union
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union
 
 from attrs import define, field