io_scene_fbx/parse_fbx.py

   1 # SPDX-FileCopyrightText: 2006-2012 assimp team
   2 # SPDX-FileCopyrightText: 2013 Blender Foundation
   3 #
   4 # SPDX-License-Identifier: GPL-2.0-or-later
   5
   6 __all__ = (
   7     "parse",
   8     "data_types",
   9     "parse_version",
  10     "FBXElem",
  11     )
  12
  13 from struct import unpack
  14 import array
  15 import zlib
  16 from io import BytesIO
  17
  18 from . import data_types
  19
  20 # at the end of each nested block, there is a NUL record to indicate
  21 # that the sub-scope exists (i.e. to distinguish between P: and P : {})
  22 _BLOCK_SENTINEL_LENGTH = ...
  23 _BLOCK_SENTINEL_DATA = ...
  24 read_fbx_elem_start = ...
  25 _IS_BIG_ENDIAN = (__import__("sys").byteorder != 'little')
  26 _HEAD_MAGIC = b'Kaydara FBX Binary\x20\x20\x00\x1a\x00'
  27 from collections import namedtuple
  28 FBXElem = namedtuple("FBXElem", ("id", "props", "props_type", "elems"))
  29 del namedtuple
  30
  31
  32 def read_uint(read):
  33     return unpack(b'<I', read(4))[0]
  34
  35
  36 def read_ubyte(read):
  37     return unpack(b'B', read(1))[0]
  38
  39
  40 def read_string_ubyte(read):
  41     size = read_ubyte(read)
  42     data = read(size)
  43     return data
  44
  45
  46 def read_array_params(read):
  47     return unpack(b'<III', read(12))
  48
  49
  50 def read_elem_start32(read):
  51     end_offset, prop_count, _prop_length, elem_id_size = unpack(b'<IIIB', read(13))
  52     elem_id = read(elem_id_size) if elem_id_size else b""
  53     return end_offset, prop_count, elem_id
  54
  55
  56 def read_elem_start64(read):
  57     end_offset, prop_count, _prop_length, elem_id_size = unpack(b'<QQQB', read(25))
  58     elem_id = read(elem_id_size) if elem_id_size else b""
  59     return end_offset, prop_count, elem_id
  60
  61
  62 def unpack_array(read, array_type, array_stride, array_byteswap):
  63     length, encoding, comp_len = read_array_params(read)
  64
  65     data = read(comp_len)
  66
  67     if encoding == 0:
  68         pass
  69     elif encoding == 1:
  70         data = zlib.decompress(data)
  71
  72     assert(length * array_stride == len(data))
  73
  74     data_array = array.array(array_type, data)
  75     if array_byteswap and _IS_BIG_ENDIAN:
  76         data_array.byteswap()
  77     return data_array
  78
  79
  80 read_data_dict = {
  81     b'Z'[0]: lambda read: unpack(b'<b', read(1))[0],  # byte
  82     b'Y'[0]: lambda read: unpack(b'<h', read(2))[0],  # 16 bit int
  83     b'B'[0]: lambda read: unpack(b'?', read(1))[0],   # 1 bit bool (yes/no)
  84     b'C'[0]: lambda read: unpack(b'<c', read(1))[0],  # char
  85     b'I'[0]: lambda read: unpack(b'<i', read(4))[0],  # 32 bit int
  86     b'F'[0]: lambda read: unpack(b'<f', read(4))[0],  # 32 bit float
  87     b'D'[0]: lambda read: unpack(b'<d', read(8))[0],  # 64 bit float
  88     b'L'[0]: lambda read: unpack(b'<q', read(8))[0],  # 64 bit int
  89     b'R'[0]: lambda read: read(read_uint(read)),      # binary data
  90     b'S'[0]: lambda read: read(read_uint(read)),      # string data
  91     b'f'[0]: lambda read: unpack_array(read, data_types.ARRAY_FLOAT32, 4, False),  # array (float)
  92     b'i'[0]: lambda read: unpack_array(read, data_types.ARRAY_INT32, 4, True),   # array (int)
  93     b'd'[0]: lambda read: unpack_array(read, data_types.ARRAY_FLOAT64, 8, False),  # array (double)
  94     b'l'[0]: lambda read: unpack_array(read, data_types.ARRAY_INT64, 8, True),   # array (long)
  95     b'b'[0]: lambda read: unpack_array(read, data_types.ARRAY_BOOL, 1, False),  # array (bool)
  96     b'c'[0]: lambda read: unpack_array(read, data_types.ARRAY_BYTE, 1, False),  # array (ubyte)
  97     }
  98
  99
 100 # FBX 7500 (aka FBX2016) introduces incompatible changes at binary level:
 101 #   * The NULL block marking end of nested stuff switches from 13 bytes long to 25 bytes long.
 102 #   * The FBX element metadata (end_offset, prop_count and prop_length) switch from uint32 to uint64.
 103 def init_version(fbx_version):
 104     global _BLOCK_SENTINEL_LENGTH, _BLOCK_SENTINEL_DATA, read_fbx_elem_start
 105
 106     _BLOCK_SENTINEL_LENGTH = ...
 107     _BLOCK_SENTINEL_DATA = ...
 108
 109     if fbx_version < 7500:
 110         _BLOCK_SENTINEL_LENGTH = 13
 111         read_fbx_elem_start = read_elem_start32
 112     else:
 113         _BLOCK_SENTINEL_LENGTH = 25
 114         read_fbx_elem_start = read_elem_start64
 115     _BLOCK_SENTINEL_DATA = (b'\0' * _BLOCK_SENTINEL_LENGTH)
 116
 117
 118 def read_elem(read, tell, use_namedtuple, tell_file_offset=0):
 119     # [0] the offset at which this block ends
 120     # [1] the number of properties in the scope
 121     # [2] the length of the property list
 122     # [3] elem name length
 123     # [4] elem name of the scope/key
 124     # read_fbx_elem_start does not return [2] because we don't use it and does not return [3] because it is only used to
 125     # get [4].
 126     end_offset, prop_count, elem_id = read_fbx_elem_start(read)
 127     if end_offset == 0:
 128         return None
 129
 130     elem_props_type = bytearray(prop_count)  # elem property types
 131     elem_props_data = [None] * prop_count    # elem properties (if any)
 132     elem_subtree = []                        # elem children (if any)
 133
 134     for i in range(prop_count):
 135         data_type = read(1)[0]
 136         elem_props_data[i] = read_data_dict[data_type](read)
 137         elem_props_type[i] = data_type
 138
 139     pos = tell()
 140     local_end_offset = end_offset - tell_file_offset
 141
 142     if pos < local_end_offset:
 143         # The default BufferedReader used when `open()`-ing files in 'rb' mode has to get the raw stream position from
 144         # the OS every time its tell() function is called. This is about 10 times slower than the tell() function of
 145         # BytesIO objects, so reading chunks of bytes from the file into memory at once and exposing them through
 146         # BytesIO can give better performance. We know the total size of each element's subtree so can read entire
 147         # subtrees into memory at a time.
 148         # The "Objects" element's subtree, however, usually makes up most of the file, so we specifically avoid reading
 149         # all its sub-elements into memory at once to reduce memory requirements at the cost of slightly worse
 150         # performance when memory is not a concern.
 151         # If we're currently reading directly from the opened file, then tell_file_offset will be zero.
 152         if tell_file_offset == 0 and elem_id != b"Objects":
 153             block_bytes_remaining = local_end_offset - pos
 154
 155             # Read the entire subtree
 156             sub_elem_bytes = read(block_bytes_remaining)
 157             num_bytes_read = len(sub_elem_bytes)
 158             if num_bytes_read != block_bytes_remaining:
 159                 raise IOError("failed to read complete nested block, expected %i bytes, but only got %i"
 160                               % (block_bytes_remaining, num_bytes_read))
 161
 162             # BytesIO provides IO API for reading bytes in memory, so we can use the same code as reading bytes directly
 163             # from a file.
 164             f = BytesIO(sub_elem_bytes)
 165             tell = f.tell
 166             read = f.read
 167             # The new `tell` function starts at zero and is offset by `pos` bytes from the start of the file.
 168             start_sub_pos = 0
 169             tell_file_offset = pos
 170             sub_tree_end = block_bytes_remaining - _BLOCK_SENTINEL_LENGTH
 171         else:
 172             # The `tell` function is unchanged, so starts at the value returned by `tell()`, which is still `pos`
 173             # because no reads have been made since then.
 174             start_sub_pos = pos
 175             sub_tree_end = local_end_offset - _BLOCK_SENTINEL_LENGTH
 176
 177         sub_pos = start_sub_pos
 178         while sub_pos < sub_tree_end:
 179             elem_subtree.append(read_elem(read, tell, use_namedtuple, tell_file_offset))
 180             sub_pos = tell()
 181
 182         # At the end of each subtree there should be a sentinel (an empty element with all bytes set to zero).
 183         if read(_BLOCK_SENTINEL_LENGTH) != _BLOCK_SENTINEL_DATA:
 184             raise IOError("failed to read nested block sentinel, "
 185                           "expected all bytes to be 0")
 186
 187         # Update `pos` for the number of bytes that have been read.
 188         pos += (sub_pos - start_sub_pos) + _BLOCK_SENTINEL_LENGTH
 189
 190     if pos != local_end_offset:
 191         raise IOError("scope length not reached, something is wrong")
 192
 193     args = (elem_id, elem_props_data, elem_props_type, elem_subtree)
 194     return FBXElem(*args) if use_namedtuple else args
 195
 196
 197 def parse_version(fn):
 198     """
 199     Return the FBX version,
 200     if the file isn't a binary FBX return zero.
 201     """
 202     with open(fn, 'rb') as f:
 203         read = f.read
 204
 205         if read(len(_HEAD_MAGIC)) != _HEAD_MAGIC:
 206             return 0
 207
 208         return read_uint(read)
 209
 210
 211 def parse(fn, use_namedtuple=True):
 212     root_elems = []
 213
 214     with open(fn, 'rb') as f:
 215         read = f.read
 216         tell = f.tell
 217
 218         if read(len(_HEAD_MAGIC)) != _HEAD_MAGIC:
 219             raise IOError("Invalid header")
 220
 221         fbx_version = read_uint(read)
 222         init_version(fbx_version)
 223
 224         while True:
 225             elem = read_elem(read, tell, use_namedtuple)
 226             if elem is None:
 227                 break
 228             root_elems.append(elem)
 229
 230     args = (b'', [], bytearray(0), root_elems)
 231     return FBXElem(*args) if use_namedtuple else args, fbx_version