From b3101abcce967c1a623c9e732199b69140a210c0 Mon Sep 17 00:00:00 2001
From: Campbell Barton <campbell@blender.org>
Date: Tue, 7 Jun 2022 11:51:53 +1000
Subject: blend_render_info: Zstd support, skip redundant file reading &
 cleanup

- Use a context manager to handle file handlers (closing both in the
  case of compressed files).

- Seek past BHead data instead of continually reading
  (checking for 'REND').

- Write errors to the stderr (so callers can differentiate it from the
  stdout).

- Use `surrogateescape` in the unlikely event of encoding errors
  so the result is always a string (possible with files pre 2.4x).

- Remove '.blend' extension check as it excludes `.blend1` files
  (we can assume the caller is passing in blend files).

- Define `__all__` to make it clear only one function is intended
  to be used.
---
 release/scripts/modules/blend_render_info.py | 114 +++++++++++++++++++--------
 1 file changed, 80 insertions(+), 34 deletions(-)

diff --git a/release/scripts/modules/blend_render_info.py b/release/scripts/modules/blend_render_info.py
index 461ea877383..4ba818a19e5 100755
--- a/release/scripts/modules/blend_render_info.py
+++ b/release/scripts/modules/blend_render_info.py
@@ -12,24 +12,65 @@
 #     int SDNAnr, nr;
 # } BHead;
 
+__all__ = (
+    "read_blend_rend_chunk",
+)
+
+
+class RawBlendFileReader:
+    """
+    Return a file handle to the raw blend file data (abstracting compressed formats).
+    """
+    __slots__ = (
+        # The path to load.
+        "_filepath",
+        # The file base file handler or None (only set for compressed formats).
+        "_blendfile_base",
+        # The file handler to return to the caller (always uncompressed data).
+        "_blendfile",
+    )
+
+    def __init__(self, filepath):
+        self._filepath = filepath
+        self._blendfile_base = None
+        self._blendfile = None
+
+    def __enter__(self):
+        blendfile = open(self._filepath, "rb")
+        blendfile_base = None
+        head = blendfile.read(4)
+        blendfile.seek(0)
+        if head[0:2] == b'\x1f\x8b':  # GZIP magic.
+            import gzip
+            blendfile_base = blendfile
+            blendfile = gzip.open(blendfile, "rb")
+        elif head[0:4] == b'\x28\xb5\x2f\xfd':  # Z-standard magic.
+            import zstandard
+            blendfile_base = blendfile
+            blendfile = zstandard.open(blendfile, "rb")
 
-def read_blend_rend_chunk(path):
+        self._blendfile_base = blendfile_base
+        self._blendfile = blendfile
 
-    import struct
+        return self._blendfile
 
-    blendfile = open(path, "rb")
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self._blendfile.close()
+        if self._blendfile_base is not None:
+            self._blendfile_base.close()
 
-    head = blendfile.read(7)
+        return False
 
-    if head[0:2] == b'\x1f\x8b':  # gzip magic
-        import gzip
-        blendfile.seek(0)
-        blendfile = gzip.open(blendfile, "rb")
-        head = blendfile.read(7)
 
+def _read_blend_rend_chunk_from_file(blendfile, filepath):
+    import struct
+    import sys
+
+    from os import SEEK_CUR
+
+    head = blendfile.read(7)
     if head != b'BLENDER':
-        print("not a blend file:", path)
-        blendfile.close()
+        sys.stderr.write("Not a blend file: %s\n" % filepath)
         return []
 
     is_64_bit = (blendfile.read(1) == b'-')
@@ -37,47 +78,52 @@ def read_blend_rend_chunk(path):
     # true for PPC, false for X86
     is_big_endian = (blendfile.read(1) == b'V')
 
-    # Now read the bhead chunk!!!
-    blendfile.read(3)  # skip the version
+    # Now read the bhead chunk!
+    blendfile.seek(3, SEEK_CUR)  # Skip the version.
 
     scenes = []
 
     sizeof_bhead = 24 if is_64_bit else 20
 
-    while blendfile.read(4) == b'REND':
-        sizeof_bhead_left = sizeof_bhead - 4
+    while len(bhead_id := blendfile.read(4)) == 4:
+        sizeof_data_left = struct.unpack('>i' if is_big_endian else '<i', blendfile.read(4))[0]
+        # 4 from the `head_id`, another 4 for the size of the BHEAD.
+        sizeof_bhead_left = sizeof_bhead - 8
 
-        struct.unpack('>i' if is_big_endian else '<i', blendfile.read(4))[0]
-        sizeof_bhead_left -= 4
+        # The remainder of the BHEAD struct is not used.
+        blendfile.seek(sizeof_bhead_left, SEEK_CUR)
 
-        # We don't care about the rest of the bhead struct
-        blendfile.read(sizeof_bhead_left)
+        if bhead_id == b'REND':
+            # Now we want the scene name, start and end frame. this is 32bits long.
+            start_frame, end_frame = struct.unpack('>2i' if is_big_endian else '<2i', blendfile.read(8))
+            sizeof_data_left -= 8
 
-        # Now we want the scene name, start and end frame. this is 32bites long
-        start_frame, end_frame = struct.unpack('>2i' if is_big_endian else '<2i', blendfile.read(8))
+            scene_name = blendfile.read(64)
+            sizeof_data_left -= 64
 
-        scene_name = blendfile.read(64)
+            scene_name = scene_name[:scene_name.index(b'\0')]
+            # It's possible old blend files are not UTF8 compliant, use `surrogateescape`.
+            scene_name = scene_name.decode("utf8", errors='surrogateescape')
 
-        scene_name = scene_name[:scene_name.index(b'\0')]
+            scenes.append((start_frame, end_frame, scene_name))
 
-        try:
-            scene_name = str(scene_name, "utf8")
-        except TypeError:
-            pass
+        if sizeof_data_left != 0:
+            blendfile.seek(sizeof_data_left, SEEK_CUR)
 
-        scenes.append((start_frame, end_frame, scene_name))
+    return scenes
 
-    blendfile.close()
 
-    return scenes
+def read_blend_rend_chunk(filepath):
+    with RawBlendFileReader(filepath) as blendfile:
+        return _read_blend_rend_chunk_from_file(blendfile, filepath)
 
 
 def main():
     import sys
-    for arg in sys.argv[1:]:
-        if arg.lower().endswith('.blend'):
-            for value in read_blend_rend_chunk(arg):
-                print("%d %d %s" % value)
+
+    for filepath in sys.argv[1:]:
+        for value in read_blend_rend_chunk(filepath):
+            print("%d %d %s" % value)
 
 
 if __name__ == '__main__':
-- 
cgit v1.2.3