From 09a0866c144e99ad012fb7ae5372cd49079047ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Wed, 7 Mar 2018 17:13:47 +0100
Subject: [PATCH] Cache open blend files

This simplifies blend file handling, ensuring that blend files are only
opened once. Otherwise it would be harder to handle things like dependency
diamonds (libraries that are referenced via multiple paths through multiple
other libraries).
---
 blender_asset_tracer/blendfile/__init__.py | 26 ++++++++++++++++++++++
 blender_asset_tracer/tracer/__init__.py    |  6 ++---
 blender_asset_tracer/tracer/file2blocks.py |  5 ++---
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py
index 1b5a8b3..bf3b61c 100644
--- a/blender_asset_tracer/blendfile/__init__.py
+++ b/blender_asset_tracer/blendfile/__init__.py
@@ -20,6 +20,7 @@
 # (c) 2014, Blender Foundation - Campbell Barton
 # (c) 2018, Blender Foundation - Sybren A. Stüvel
 
+import atexit
 import collections
 import gzip
 import logging
@@ -27,6 +28,8 @@ import os
 import struct
 import pathlib
 import tempfile
+
+import functools
 import typing
 
 from . import exceptions, dna_io, dna, header
@@ -39,6 +42,29 @@ FILE_BUFFER_SIZE = 1024 * 1024
 BLENDFILE_MAGIC = b'BLENDER'
 GZIP_MAGIC = b'\x1f\x8b'
 
+_cached_bfiles = {}
+
+
+def open_cached(path: pathlib.Path, mode='rb') -> 'BlendFile':
+    """Open a blend file, ensuring it is only opened once."""
+    bfile_path = path.absolute().resolve()
+    try:
+        return _cached_bfiles[bfile_path]
+    except KeyError:
+        pass
+
+    bfile = BlendFile(path, mode=mode)
+    _cached_bfiles[bfile_path] = bfile
+    return bfile
+
+
+@atexit.register
+def close_all_cached():
+    log.info('Closing all blend files')
+    for bfile in _cached_bfiles.values():
+        bfile.close()
+    _cached_bfiles.clear()
+
 
 class BlendFile:
     """Representation of a blend file.
diff --git a/blender_asset_tracer/tracer/__init__.py b/blender_asset_tracer/tracer/__init__.py
index 366aec7..f27ad57 100644
--- a/blender_asset_tracer/tracer/__init__.py
+++ b/blender_asset_tracer/tracer/__init__.py
@@ -23,9 +23,9 @@ def deps(bfilepath: pathlib.Path) -> typing.Iterator[result.BlockUsage]:
     :param bfilepath: File to open.
     """
 
-    with blendfile.BlendFile(bfilepath) as bfile:
-        for block in asset_holding_blocks(file2blocks.iter_blocks(bfile)):
-            yield from blocks2assets.iter_assets(block)
+    bfile = blendfile.open_cached(bfilepath)
+    for block in asset_holding_blocks(file2blocks.iter_blocks(bfile)):
+        yield from blocks2assets.iter_assets(block)
 
 
 def asset_holding_blocks(blocks: typing.Iterable[blendfile.BlendFileBlock]) \
diff --git a/blender_asset_tracer/tracer/file2blocks.py b/blender_asset_tracer/tracer/file2blocks.py
index 283f720..a9b7bba 100644
--- a/blender_asset_tracer/tracer/file2blocks.py
+++ b/blender_asset_tracer/tracer/file2blocks.py
@@ -86,9 +86,8 @@ class _BlockIterator:
                 continue
 
             log.debug('Expanding %d blocks in %s', len(idblocks), lib_path)
-
-            with blendfile.BlendFile(lib_path) as libfile:
-                yield from self.iter_blocks(libfile, idblocks)
+            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)