Ensure module caching follows runtime creation

src/Bonsai.Scripting.Python/Eval.cs

@@ -44,9 +44,11 @@ public class Eval : Combinator<PyObject>
         /// </returns>
         public override IObservable<PyObject> Process<TSource>(IObservable<TSource> source)
         {
-            return RuntimeManager.RuntimeSource
-                .GetModuleOrDefaultAsync(Module)
-                .SelectMany(module => source.Select(_ => module.Eval(Expression)));
+            return RuntimeManager.RuntimeSource.SelectMany(runtime =>
+            {
+                var module = Module ?? runtime.MainModule;
+                return source.Select(_ => module.Eval(Expression));
+            });
         }
 
         /// <summary>

src/Bonsai.Scripting.Python/Exec.cs

@@ -45,9 +45,11 @@ public class Exec : Combinator<PyModule>
         /// </returns>
         public override IObservable<PyModule> Process<TSource>(IObservable<TSource> source)
         {
-            return RuntimeManager.RuntimeSource
-                .GetModuleOrDefaultAsync(Module)
-                .SelectMany(module => source.Select(_ => module.Exec(Script)));
+            return RuntimeManager.RuntimeSource.SelectMany(runtime =>
+            {
+                var module = Module ?? runtime.MainModule;
+                return source.Select(_ => module.Exec(Script));
+            });
         }
 
         /// <summary>

src/Bonsai.Scripting.Python/Get.cs

@@ -38,10 +38,9 @@ public class Get : Source<PyObject>
         /// </returns>
         public override IObservable<PyObject> Generate()
         {
-            return RuntimeManager.RuntimeSource
-                .GetModuleOrDefaultAsync(Module)
-                .ObserveOnGIL()
-                .Select(module => module.Get(VariableName));
+            return from runtime in RuntimeManager.RuntimeSource.ObserveOnGIL()
+                   let module = Module ?? runtime.MainModule
+                   select module.Get(VariableName);
         }
 
         /// <summary>
@@ -60,9 +59,11 @@ public override IObservable<PyObject> Generate()
         /// </returns>
         public IObservable<PyObject> Generate<TSource>(IObservable<TSource> source)
         {
-            return RuntimeManager.RuntimeSource
-                .GetModuleOrDefaultAsync(Module)
-                .SelectMany(module => source.Select(_ => module.Get(VariableName)));
+            return RuntimeManager.RuntimeSource.SelectMany(runtime =>
+            {
+                var module = Module ?? runtime.MainModule;
+                return source.Select(_ => module.Get(VariableName));
+            });
         }
 
         /// <summary>

src/Bonsai.Scripting.Python/ObservableExtensions.cs

@@ -7,13 +7,6 @@ namespace Bonsai.Scripting.Python
 {
     static class ObservableExtensions
     {
-        public static IObservable<PyModule> GetModuleOrDefaultAsync(this IObservable<RuntimeManager> source, PyModule module)
-        {
-            return module != null
-                ? Observable.Return(module)
-                : source.Select(runtime => runtime.MainModule);
-        }
-
         public static IObservable<TSource> ObserveOnGIL<TSource>(this IObservable<TSource> source)
         {
             return Observable.Create<TSource>(observer =>

src/Bonsai.Scripting.Python/Set.cs

@@ -45,9 +45,11 @@ public class Set : Sink
         /// </returns>
         public override IObservable<TSource> Process<TSource>(IObservable<TSource> source)
         {
-            return RuntimeManager.RuntimeSource
-                .GetModuleOrDefaultAsync(Module)
-                .SelectMany(module => source.Do(value => module.Set(VariableName, value)));
+            return RuntimeManager.RuntimeSource.SelectMany(runtime =>
+            {
+                var module = Module ?? runtime.MainModule;
+                return source.Do(value => module.Set(VariableName, value));
+            });
         }
     }
 }