From 03615c30b566eb31033854bee6f0e7e9d8d4b5ba Mon Sep 17 00:00:00 2001
From: Adam Johnson <AdamJohnso@gmail.com>
Date: Sun, 11 Jun 2023 18:34:51 -0400
Subject: [PATCH] Improve `pre_export()` generators.

It is now possible to write
```python
my_object = yield create_something()
```
to ensure that temporary objects are always managed by the exporter. Previously, the more verbose
```python
my_object = create_something()
yield my_object
```
was needed.
---
 korman/exporter/convert.py | 42 +++++++++++++++++++++++++++++++++-----
 1 file changed, 37 insertions(+), 5 deletions(-)

diff --git a/korman/exporter/convert.py b/korman/exporter/convert.py
index 17588c3..c320c6a 100644
--- a/korman/exporter/convert.py
+++ b/korman/exporter/convert.py
@@ -18,6 +18,7 @@ import bpy
 from collections import defaultdict
 from contextlib import ExitStack
 import functools
+import inspect
 from pathlib import Path
 from typing import *
 
@@ -459,10 +460,14 @@ class Exporter:
             if hasattr(temporary, "__enter__"):
                 ctx_temporary = self.exit_stack.enter_context(temporary)
                 if ctx_temporary is not None:
-                    handle_temporary(ctx_temporary, parent)
+                    return handle_temporary(ctx_temporary, parent)
             else:
                 raise RuntimeError("Temporary object of type '{}' generated by '{}' was unhandled".format(temporary.__class__, parent.name))
 
+        @handle_temporary.register(utils.BMeshObject)
+        def _(temporary, parent):
+            return handle_temporary(temporary.release(), parent)
+
         @handle_temporary.register(bpy.types.Object)
         def _(temporary, parent):
             self.exit_stack.enter_context(TemporaryObject(temporary, bpy.data.objects.remove))
@@ -483,6 +488,7 @@ class Exporter:
                     for mod in temporary.plasma_modifiers.modifiers:
                         mod.sanity_check()
                     do_pre_export(temporary)
+            return temporary
 
         @handle_temporary.register(bpy.types.NodeTree)
         def _(temporary, parent):
@@ -491,15 +497,41 @@ class Exporter:
             if temporary.bl_idname == "PlasmaNodeTree":
                 parent_so = self.mgr.find_create_object(plSceneObject, bl=parent)
                 self.want_node_trees[temporary.name].add((parent, parent_so))
+            return temporary
 
         def do_pre_export(bo):
             for mod in bo.plasma_modifiers.modifiers:
                 proc = getattr(mod, "pre_export", None)
                 if proc is not None:
-                    result = proc(self, bo)
-                    if result is not None:
-                        for i in filter(None, result):
-                            handle_temporary(i, bo)
+                    # pre_export() should be a *generator*. Generators are bidirectional, although
+                    # trivial usages simply yield values for the consumer. What we want to do here
+                    # is use the bidirectional nature of generators to simplify the code in pre_export().
+                    # We will pump the pre_export generator. With each pump, we will send the generator
+                    # the very thing it gave us. That way, if pre_export needs to do some work on
+                    # the object it generated, then it can do:
+                    # ```
+                    # my_object = yield create_something()
+                    # my_object.foo = bar
+                    # ```
+                    # instead of the more verbose
+                    # ```
+                    # my_object = create_something()
+                    # yield my_object
+                    # my_object.foo = bar
+                    # ```
+                    assert inspect.isgeneratorfunction(proc), "pre_export should be a generator function"
+                    pre_result = proc(self, bo)
+                    try:
+                        gen_result = None
+                        while True:
+                            gen_result = pre_result.send(gen_result)
+                            if gen_result is not None:
+                                gen_result = handle_temporary(gen_result, bo)
+                    except StopIteration as e:
+                        if e.value is not None:
+                            handle_temporary(e.value, bo)
+                    finally:
+                        pre_result.close()
 
         with indent():
             for bl_obj in self._objects: