Optimize creating AssetManager further

src/retro_data_structures/asset_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import collections
 import fnmatch
 import json
 import logging
@@ -146,13 +147,9 @@ def _resolve_asset_id(self, value: NameOrAssetId) -> AssetId:
  return self._custom_asset_ids[str(value)]
  return resolve_asset_id(self.target_game, value)
 
- def _add_pak_name_for_asset_id(self, asset_id: AssetId, pak_name: str):
- self._paks_for_asset_id[asset_id] = self._paks_for_asset_id.get(asset_id, set())
- self._paks_for_asset_id[asset_id].add(pak_name)
-
  def _update_headers(self):
  self._ensured_asset_ids = {}
- self._paks_for_asset_id = {}
+ self._paks_for_asset_id = collections.defaultdict(set)
  self._types_for_asset_id = {}
 
  self._custom_asset_ids = {}
@@ -170,8 +167,8 @@ def _update_headers(self):
 
  self._ensured_asset_ids[name] = set()
  for entry in pak_no_data.resources:
- self._add_pak_name_for_asset_id(entry.asset.id, name)
- self._types_for_asset_id[entry.asset.id] = entry.asset.type
+ self._paks_for_asset_id[entry.asset_id].add(name)
+ self._types_for_asset_id[entry.asset_id] = entry.asset_type
 
  def all_asset_ids(self) -> Iterator[AssetId]:
  """

src/retro_data_structures/formats/pak_gc.py

@@ -8,7 +8,7 @@
 
 from retro_data_structures import game_check
 from retro_data_structures.base_resource import AssetId, AssetType, Dependency
-from retro_data_structures.common_types import ObjectTag_32
+from retro_data_structures.common_types import AssetId32, FourCC
 from retro_data_structures.compression import LZOCompressedBlock, ZlibCompressedBlock
 from retro_data_structures.construct_extensions.alignment import AlignTo
 
@@ -20,23 +20,41 @@
  version_minor=Const(5, Int16ub),
  unused=Const(0, Int32ub),
 )
-ResourceHeader = Struct(
+ConstructResourceHeader = Struct(
  compressed=Int32ub,
- asset=ObjectTag_32,
+ asset_type=FourCC,
+ asset_id=AssetId32,
  size=Int32ub,
  offset=Int32ub,
 )
+
+
+def _emitparse_header(code: construct.CodeGen) -> str:
+ code.append("ResourceHeader_Format = struct.Struct('>LLLLL')")
+ code.append(
+ """def _create_resource_header(compressed, asset_type, asset_id, size, offset):
+ return Container(compressed=compressed, asset_type=asset_type.to_bytes(4, "big").decode("ascii"),
+ asset_id=asset_id, size=size, offset=offset)
+ """
+ )
+ return "_create_resource_header(*ResourceHeader_Format.unpack(io.read(20)))"
+
+
+ConstructResourceHeader._emitparse = _emitparse_header
+
 PAKNoData = Struct(
  _header=PAKHeader,
  named_resources=PrefixedArray(
  Int32ub,
  Struct(
- asset=ObjectTag_32,
+ asset_type=FourCC,
+ asset_id=AssetId32,
  name=PascalString(Int32ub, "utf-8"),
  ),
  ),
- resources=PrefixedArray(Int32ub, ResourceHeader),
+ resources=PrefixedArray(Int32ub, ConstructResourceHeader),
 ).compile()
+
 CompressedPakResource = FocusedSeq(
  "data",
  decompressed_size=Rebuild(Int32ub, construct.len_(construct.this.data)),
@@ -108,8 +126,8 @@ def _parse(self, stream, context, path) -> PakBody:
 
  files.append(
  PakFile(
- resource.asset.id,
- resource.asset.type,
+ resource.asset_id,
+ resource.asset_type,
  resource.compressed > 0,
  uncompressed_data,
  compressed_data,
@@ -118,7 +136,7 @@ def _parse(self, stream, context, path) -> PakBody:
 
  return PakBody(
  named_resources={
- named.name: Dependency(type=named.asset.type, id=named.asset.id) for named in header.named_resources
+ named.name: Dependency(type=named.asset_type, id=named.asset_id) for named in header.named_resources
  },
  files=files,
  )
@@ -130,21 +148,17 @@ def _build(self, obj: PakBody, stream, context, path):
  _header=construct.Container(),
  named_resources=construct.ListContainer(
  construct.Container(
- asset=construct.Container(
- type=dep.type,
- id=dep.id,
- ),
+ asset_type=dep.type,
+ asset_id=dep.id,
  name=name,
  )
  for name, dep in obj.named_resources.items()
  ),
  resources=construct.ListContainer(
  construct.Container(
  compressed=0,
- asset=construct.Container(
- type=file.asset_type,
- id=file.asset_id,
- ),
+ asset_type=file.asset_type,
+ asset_id=file.asset_id,
  size=0,
  offset=0,
  )

src/retro_data_structures/formats/pak_wiiu.py

@@ -13,10 +13,8 @@
 from retro_data_structures.formats.pak_gc import PakFile
 
 StringTableEntry = Struct(
- asset=Struct(
- type=FourCC,
- id=GUID,
- ),
+ asset_type=FourCC,
+ asset_id=GUID,
  name=construct.PascalString(Int32ul, "utf8"),
 )
 
@@ -26,10 +24,8 @@
 )
 
 AssetDirectoryEntry = Struct(
- asset=Struct(
- type=FourCC,
- id=GUID,
- ),
+ asset_type=FourCC,
+ asset_id=GUID,
  version_a=Hex(Int32ul),
  version_b=Hex(Int32ul),
  offset=Hex(Int64ul),
@@ -87,14 +83,14 @@ def _parse(self, stream, context, path):
  for i, resource in enumerate(header.tocc.ADIR.data):
  # if resource.decompressed_size != resource.size:
  # raise construct.ConstructError(
- # f"Resource {i} ({resource.asset.id} {resource.asset.type}), is compressed", path)
+ # f"Resource {i} ({resource.asset_id} {resource.asset_type}), is compressed", path)
 
  construct.stream_seek(stream, resource.offset, 0, path)
  data = construct.stream_read(stream, resource.size, path)
  files.append(
  PakFile(
- resource.asset.id,
- resource.asset.type,
+ resource.asset_id,
+ resource.asset_type,
  False,
  data,
  None,
@@ -126,10 +122,8 @@ def _build(self, obj, stream, context, path):
  id="ADIR",
  data=construct.ListContainer(
  construct.Container(
- asset=construct.Container(
- type=file.asset_type,
- id=file.asset_id,
- ),
+ asset_type=file.asset_type,
+ asset_id=file.asset_id,
  version_a=file.extra.version_a,
  version_b=file.extra.version_b,
  offset=0,

src/retro_data_structures/formats/room.py

@@ -6,7 +6,6 @@
 import construct
 from construct import Container, Hex, Int16ul, Int32ul, PrefixedArray, Struct
 
-import retro_data_structures.properties.prime_remastered.objects
 from retro_data_structures.base_resource import AssetType, BaseResource, Dependency
 from retro_data_structures.common_types import GUID
 from retro_data_structures.construct_extensions.misc import ErrorWithMessage, UntilEof
@@ -193,6 +192,8 @@
  self.type_id = type_id
 
  def _decode(self, obj: bytes, context, path):
+ import retro_data_structures.properties.prime_remastered.objects
+
  type_id = self.type_id(context)
  try:
  property_class = retro_data_structures.properties.prime_remastered.objects.get_object(type_id)
@@ -208,6 +209,8 @@
 
 class ComponentTypeName(int):
  def __str__(self) -> str:
+ import retro_data_structures.properties.prime_remastered.objects
+
  try:
  return retro_data_structures.properties.prime_remastered.objects.get_object(self).__name__
  except KeyError: