177 lines
6.9 KiB
Python
177 lines
6.9 KiB
Python
|
# ***** BEGIN GPL LICENSE BLOCK *****
|
||
|
#
|
||
|
# This program is free software; you can redistribute it and/or
|
||
|
# modify it under the terms of the GNU General Public License
|
||
|
# as published by the Free Software Foundation; either version 2
|
||
|
# of the License, or (at your option) any later version.
|
||
|
#
|
||
|
# This program is distributed in the hope that it will be useful,
|
||
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||
|
# GNU General Public License for more details.
|
||
|
#
|
||
|
# You should have received a copy of the GNU General Public License
|
||
|
# along with this program; if not, write to the Free Software Foundation,
|
||
|
# Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
|
||
|
#
|
||
|
# ***** END GPL LICENCE BLOCK *****
|
||
|
#
|
||
|
# (c) 2018, Blender Foundation - Sybren A. Stüvel
|
||
|
"""Expand data blocks.
|
||
|
|
||
|
The expansion process follows pointers and library links to construct the full
|
||
|
set of actually-used data blocks. This set consists of all data blocks in the
|
||
|
initial blend file, and all *actually linked-to* data blocks in linked
|
||
|
blend files.
|
||
|
"""
|
||
|
import collections
|
||
|
import logging
|
||
|
import pathlib
|
||
|
import queue
|
||
|
import typing
|
||
|
|
||
|
from blender_asset_tracer import blendfile, bpathlib
|
||
|
from . import expanders, progress
|
||
|
|
||
|
_funcs_for_code = {} # type: typing.Dict[bytes, typing.Callable]
|
||
|
log = logging.getLogger(__name__)
|
||
|
|
||
|
|
||
|
# noinspection PyProtectedMember
|
||
|
class BlockQueue(queue.PriorityQueue):
|
||
|
"""PriorityQueue that sorts by filepath and file offset"""
|
||
|
|
||
|
def _put(self, item: blendfile.BlendFileBlock):
|
||
|
super()._put((item.bfile.filepath, item.file_offset, item))
|
||
|
|
||
|
def _get(self) -> blendfile.BlendFileBlock:
|
||
|
_, _, item = super()._get()
|
||
|
return item
|
||
|
|
||
|
|
||
|
class BlockIterator:
|
||
|
"""Expand blocks with dependencies from other libraries.
|
||
|
|
||
|
This class exists so that we have some context for the recursive expansion
|
||
|
without having to pass those variables to each recursive call.
|
||
|
"""
|
||
|
|
||
|
def __init__(self) -> None:
|
||
|
# Set of (blend file Path, block address) of already-reported blocks.
|
||
|
self.blocks_yielded = set() # type: typing.Set[typing.Tuple[pathlib.Path, int]]
|
||
|
|
||
|
# Queue of blocks to visit
|
||
|
self.to_visit = BlockQueue()
|
||
|
|
||
|
self.progress_cb = progress.Callback()
|
||
|
|
||
|
def iter_blocks(self,
|
||
|
bfile: blendfile.BlendFile,
|
||
|
limit_to: typing.Set[blendfile.BlendFileBlock] = set(),
|
||
|
) -> typing.Iterator[blendfile.BlendFileBlock]:
|
||
|
"""Expand blocks with dependencies from other libraries."""
|
||
|
|
||
|
self.progress_cb.trace_blendfile(bfile.filepath)
|
||
|
log.info('inspecting: %s', bfile.filepath)
|
||
|
if limit_to:
|
||
|
self._queue_named_blocks(bfile, limit_to)
|
||
|
else:
|
||
|
self._queue_all_blocks(bfile)
|
||
|
|
||
|
blocks_per_lib = yield from self._visit_blocks(bfile, limit_to)
|
||
|
yield from self._visit_linked_blocks(blocks_per_lib)
|
||
|
|
||
|
def _visit_blocks(self, bfile, limit_to):
|
||
|
bpath = bpathlib.make_absolute(bfile.filepath)
|
||
|
root_dir = bpathlib.BlendPath(bpath.parent)
|
||
|
|
||
|
# Mapping from library path to data blocks to expand.
|
||
|
blocks_per_lib = collections.defaultdict(set)
|
||
|
|
||
|
while not self.to_visit.empty():
|
||
|
block = self.to_visit.get()
|
||
|
assert isinstance(block, blendfile.BlendFileBlock)
|
||
|
if (bpath, block.addr_old) in self.blocks_yielded:
|
||
|
continue
|
||
|
|
||
|
if block.code == b'ID':
|
||
|
# ID blocks represent linked-in assets. Those are the ones that
|
||
|
# should be loaded from their own blend file and "expanded" to
|
||
|
# the entire set of data blocks required to render them. We
|
||
|
# defer the handling of those so that we can work with one
|
||
|
# blend file at a time.
|
||
|
lib = block.get_pointer(b'lib')
|
||
|
lib_bpath = bpathlib.BlendPath(lib[b'name']).absolute(root_dir)
|
||
|
blocks_per_lib[lib_bpath].add(block)
|
||
|
|
||
|
# The library block itself should also be reported, because it
|
||
|
# represents a blend file that is a dependency as well.
|
||
|
self.to_visit.put(lib)
|
||
|
continue
|
||
|
|
||
|
if limit_to:
|
||
|
# We're limiting the blocks, so we have to expand them to make
|
||
|
# sure we don't miss anything. Otherwise we're yielding the
|
||
|
# entire file anyway, and no expansion is necessary.
|
||
|
self._queue_dependencies(block)
|
||
|
self.blocks_yielded.add((bpath, block.addr_old))
|
||
|
yield block
|
||
|
|
||
|
return blocks_per_lib
|
||
|
|
||
|
def _visit_linked_blocks(self, blocks_per_lib):
|
||
|
# We've gone through all the blocks in this file, now open the libraries
|
||
|
# and iterate over the blocks referred there.
|
||
|
for lib_bpath, idblocks in blocks_per_lib.items():
|
||
|
lib_path = bpathlib.make_absolute(lib_bpath.to_path())
|
||
|
|
||
|
#assert lib_path.exists()
|
||
|
if not lib_path.exists():
|
||
|
log.warning('Library %s does not exist', lib_path)
|
||
|
continue
|
||
|
|
||
|
log.debug('Expanding %d blocks in %s', len(idblocks), lib_path)
|
||
|
libfile = blendfile.open_cached(lib_path)
|
||
|
yield from self.iter_blocks(libfile, idblocks)
|
||
|
|
||
|
def _queue_all_blocks(self, bfile: blendfile.BlendFile):
|
||
|
log.debug('Queueing all blocks from file %s', bfile.filepath)
|
||
|
for block in bfile.blocks:
|
||
|
# Don't bother visiting DATA blocks, as we won't know what
|
||
|
# to do with them anyway.
|
||
|
if block.code == b'DATA':
|
||
|
continue
|
||
|
self.to_visit.put(block)
|
||
|
|
||
|
def _queue_named_blocks(self,
|
||
|
bfile: blendfile.BlendFile,
|
||
|
limit_to: typing.Set[blendfile.BlendFileBlock]):
|
||
|
"""Queue only the blocks referred to in limit_to.
|
||
|
|
||
|
:param bfile:
|
||
|
:param limit_to: set of ID blocks that name the blocks to queue.
|
||
|
The queued blocks are loaded from the actual blend file, and
|
||
|
selected by name.
|
||
|
"""
|
||
|
|
||
|
for to_find in limit_to:
|
||
|
assert to_find.code == b'ID'
|
||
|
name_to_find = to_find[b'name']
|
||
|
code = name_to_find[:2]
|
||
|
log.debug('Finding block %r with code %r', name_to_find, code)
|
||
|
same_code = bfile.find_blocks_from_code(code)
|
||
|
for block in same_code:
|
||
|
if block.id_name == name_to_find:
|
||
|
log.debug('Queueing %r from file %s', block, bfile.filepath)
|
||
|
self.to_visit.put(block)
|
||
|
|
||
|
def _queue_dependencies(self, block: blendfile.BlendFileBlock):
|
||
|
for block in expanders.expand_block(block):
|
||
|
self.to_visit.put(block)
|
||
|
|
||
|
|
||
|
def iter_blocks(bfile: blendfile.BlendFile) -> typing.Iterator[blendfile.BlendFileBlock]:
|
||
|
"""Generator, yield all blocks in this file + required blocks in libs."""
|
||
|
bi = BlockIterator()
|
||
|
yield from bi.iter_blocks(bfile)
|