From 21fb18247dd176c66fe70036bb11b8a6e9c34b0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Mon, 26 Mar 2018 16:46:13 +0200
Subject: [PATCH] Find textures and images of various modifiers

- Wave
- UV Project
- Weight VG Edit
- Weight VG Mix
- Weight VG Proximity
---
 blender_asset_tracer/cdefs.py                 |  5 ++
 .../trace/modifier_walkers.py                 | 48 ++++++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/blender_asset_tracer/cdefs.py b/blender_asset_tracer/cdefs.py
index 2414fe3..16e1243 100644
--- a/blender_asset_tracer/cdefs.py
+++ b/blender_asset_tracer/cdefs.py
@@ -34,8 +34,13 @@ IMA_SRC_SEQUENCE = 2
 IMA_SRC_MOVIE = 3
 
 # DNA_modifier_types.h
+eModifierType_Wave = 7
 eModifierType_Displace = 14
+eModifierType_UVProject = 15
 eModifierType_ParticleSystem = 19
+eModifierType_WeightVGEdit = 36
+eModifierType_WeightVGMix = 37
+eModifierType_WeightVGProximity = 38
 eModifierType_Ocean = 39
 eModifierType_MeshCache = 46
 
diff --git a/blender_asset_tracer/trace/modifier_walkers.py b/blender_asset_tracer/trace/modifier_walkers.py
index cbcf76b..075ff5a 100644
--- a/blender_asset_tracer/trace/modifier_walkers.py
+++ b/blender_asset_tracer/trace/modifier_walkers.py
@@ -62,18 +62,54 @@ def modifier_ocean(modifier: blendfile.BlendFileBlock, block_name: bytes) \
                             block_name=block_name)
 
 
-@mod_handler(cdefs.eModifierType_Displace)
-def modifier_texture(modifier: blendfile.BlendFileBlock, block_name: bytes) \
+def _get_texture(prop_name: bytes, dblock: blendfile.BlendFileBlock, block_name: bytes) \
         -> typing.Iterator[result.BlockUsage]:
-    tx = modifier.get_pointer(b'texture')
-    if not tx:
+    """Yield block usages from a texture propery.
+
+    Assumes dblock[prop_name] is a texture data block.
+    """
+    if dblock is None:
         return
-    ima = tx.get_pointer(b'ima')
+
+    tx = dblock.get_pointer(prop_name)
+    yield from _get_image(b'ima', tx, block_name)
+
+
+def _get_image(prop_name: bytes, dblock: blendfile.BlendFileBlock, block_name: bytes) \
+        -> typing.Iterator[result.BlockUsage]:
+    """Yield block usages from an image propery.
+
+    Assumes dblock[prop_name] is an image data block.
+    """
+    if not dblock:
+        return
+    ima = dblock.get_pointer(prop_name)
     if not ima:
         return
 
     path, field = ima.get(b'name', return_field=True)
-    yield result.BlockUsage(modifier, path, path_full_field=field, block_name=block_name)
+    yield result.BlockUsage(ima, path, path_full_field=field, block_name=block_name)
+
+
+@mod_handler(cdefs.eModifierType_Displace)
+@mod_handler(cdefs.eModifierType_Wave)
+def modifier_texture(modifier: blendfile.BlendFileBlock, block_name: bytes) \
+        -> typing.Iterator[result.BlockUsage]:
+    return _get_texture(b'texture', modifier, block_name)
+
+
+@mod_handler(cdefs.eModifierType_WeightVGEdit)
+@mod_handler(cdefs.eModifierType_WeightVGMix)
+@mod_handler(cdefs.eModifierType_WeightVGProximity)
+def modifier_mask_texture(modifier: blendfile.BlendFileBlock, block_name: bytes) \
+        -> typing.Iterator[result.BlockUsage]:
+    return _get_texture(b'mask_texture', modifier, block_name)
+
+
+@mod_handler(cdefs.eModifierType_UVProject)
+def modifier_image(modifier: blendfile.BlendFileBlock, block_name: bytes) \
+        -> typing.Iterator[result.BlockUsage]:
+    yield from _get_image(b'image', modifier, block_name)
 
 
 @mod_handler(cdefs.eModifierType_ParticleSystem)