Change GroovyScriptExtension.load to use a map of context objects instead of a ComputationManager to be more generic

iidm/iidm-scripting/src/main/groovy/com/powsybl/iidm/network/scripting/NetworkLoadSaveGroovyScriptExtension.groovy

@@ -8,7 +8,6 @@
 package com.powsybl.iidm.network.scripting
 
 import com.google.auto.service.AutoService
-import com.powsybl.computation.ComputationManager
 import com.powsybl.computation.local.LocalComputationManager
 import com.powsybl.iidm.network.ExportersLoader
 import com.powsybl.iidm.network.ExportersServiceLoader
@@ -48,7 +47,7 @@ class NetworkLoadSaveGroovyScriptExtension implements GroovyScriptExtension {
     }
 
     @Override
-    void load(Binding binding, ComputationManager computationManager) {
+    void load(Binding binding, Map<Class<?>, Object> contextObjects) {
         binding.loadNetwork = { String file, Properties parameters = null ->
             Network.read(fileSystem.getPath(file), LocalComputationManager.getDefault(),
                     importConfig, parameters, importersLoader)

loadflow/loadflow-scripting/src/main/groovy/com/powsybl/loadflow/scripting/LoadFlowGroovyScriptExtension.groovy

@@ -23,21 +23,24 @@ class LoadFlowGroovyScriptExtension implements GroovyScriptExtension {
     private final LoadFlowParameters parameters
 
     LoadFlowGroovyScriptExtension(LoadFlowParameters parameters) {
-        assert parameters
-        this.parameters = parameters
+        this.parameters = Objects.requireNonNull(parameters)
     }
 
     LoadFlowGroovyScriptExtension() {
         this(LoadFlowParameters.load())
     }
 
     @Override
-    void load(Binding binding, ComputationManager computationManager) {
-        binding.loadFlow = { Network network, LoadFlowParameters parameters = this.parameters ->
-            LoadFlow.run(network, network.getVariantManager().getWorkingVariantId(), computationManager, parameters)
-        }
-        binding.loadflow = { Network network, LoadFlowParameters parameters = this.parameters ->
-            LoadFlow.run(network, network.getVariantManager().getWorkingVariantId(), computationManager, parameters)
+    void load(Binding binding, Map<Class<?>, Object> contextObjects) {
+        if (contextObjects.keySet().contains(ComputationManager.class)) {
+            ComputationManager computationManager = contextObjects.get(ComputationManager.class) as ComputationManager
+
+            binding.loadFlow = { Network network, LoadFlowParameters parameters = this.parameters ->
+                LoadFlow.run(network, network.getVariantManager().getWorkingVariantId(), computationManager, parameters)
+            }
+            binding.loadflow = { Network network, LoadFlowParameters parameters = this.parameters ->
+                LoadFlow.run(network, network.getVariantManager().getWorkingVariantId(), computationManager, parameters)
+            }
         }
     }
 

loadflow/loadflow-scripting/src/test/java/com/powsybl/loadflow/scripting/LoadFlowExtensionGroovyScriptTest.java

@@ -7,7 +7,6 @@
  */
 package com.powsybl.loadflow.scripting;
 
-import com.powsybl.computation.ComputationManager;
 import com.powsybl.iidm.network.Network;
 import com.powsybl.iidm.network.VariantManager;
 import com.powsybl.iidm.network.VariantManagerConstants;
@@ -21,14 +20,14 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 /**
  * @author Geoffroy Jamgotchian {@literal <geoffroy.jamgotchian at rte-france.com>}
  */
 class LoadFlowExtensionGroovyScriptTest extends AbstractGroovyScriptTest {
 
     private Network fooNetwork;
-    private ComputationManager computationManager;
 
     @BeforeEach
     void setUp() {
@@ -57,12 +56,13 @@ protected String getExpectedOutput() {
     protected List<GroovyScriptExtension> getExtensions() {
         GroovyScriptExtension ext = new GroovyScriptExtension() {
             @Override
-            public void load(Binding binding, ComputationManager computationManager) {
+            public void load(Binding binding, Map<Class<?>, Object> contextObjects) {
                 binding.setVariable("n", fooNetwork);
             }
 
             @Override
             public void unload() {
+                // Nothing to do here
             }
         };
 

scripting-test/src/main/java/com/powsybl/scripting/test/AbstractGroovyScriptTest.java

@@ -19,7 +19,9 @@
 import java.io.IOException;
 import java.io.StringWriter;
 import java.io.UncheckedIOException;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
@@ -35,7 +37,8 @@ public abstract class AbstractGroovyScriptTest {
     protected abstract List<GroovyScriptExtension> getExtensions();
 
     public void doTest() {
-        ComputationManager computationManager = Mockito.mock(ComputationManager.class);
+        Map<Class<?>, Object> contextObjects = new HashMap<>();
+        contextObjects.put(ComputationManager.class, Mockito.mock(ComputationManager.class));
         Binding binding = new Binding();
         StringWriter out = null;
         try {
@@ -47,7 +50,7 @@ public void doTest() {
                 // Add a check on thread interruption in every loop (for, while) in the script
                 conf.addCompilationCustomizers(new ASTTransformationCustomizer(ThreadInterrupt.class));
 
-                getExtensions().forEach(it -> it.load(binding, computationManager));
+                getExtensions().forEach(it -> it.load(binding, contextObjects));
                 GroovyShell shell = new GroovyShell(binding, conf);
                 shell.evaluate(getCode());
                 out = (StringWriter) binding.getProperty("out");

scripting/src/main/groovy/com/powsybl/scripting/groovy/GroovyScripts.groovy

@@ -7,6 +7,7 @@
  */
 package com.powsybl.scripting.groovy
 
+import com.powsybl.computation.ComputationManager
 import com.powsybl.computation.DefaultComputationManagerConfig
 import groovy.transform.ThreadInterrupt
 import org.codehaus.groovy.control.CompilerConfiguration
@@ -47,6 +48,10 @@ class GroovyScripts {
     }
 
     static void run(Reader codeReader, Binding binding, Iterable<GroovyScriptExtension> extensions, PrintStream out) {
+        run(codeReader, binding, extensions, out, new HashMap<>())
+    }
+
+    static void run(Reader codeReader, Binding binding, Iterable<GroovyScriptExtension> extensions, PrintStream out, Map<Class<?>, Objects> contextObjects) {
         assert codeReader
         assert extensions != null
 
@@ -58,14 +63,15 @@ class GroovyScripts {
         // Computation manager
         DefaultComputationManagerConfig config = DefaultComputationManagerConfig.load()
         binding.computationManager = config.createShortTimeExecutionComputationManager()
+        contextObjects.put(ComputationManager.class, binding.computationManager)
 
         if (out != null) {
             binding.out = out
         }
 
         try {
             // load extensions
-            extensions.forEach { it.load(binding, binding.computationManager) }
+            extensions.forEach { it.load(binding, contextObjects) }
 
             GroovyShell shell = new GroovyShell(binding, conf)
 

scripting/src/main/java/com/powsybl/scripting/groovy/GroovyScriptExtension.java

@@ -7,15 +7,16 @@
  */
 package com.powsybl.scripting.groovy;
 
-import com.powsybl.computation.ComputationManager;
 import groovy.lang.Binding;
 
+import java.util.Map;
+
 /**
  * @author Geoffroy Jamgotchian {@literal <geoffroy.jamgotchian at rte-france.com>}
  */
 public interface GroovyScriptExtension {
 
-    void load(Binding binding, ComputationManager computationManager);
+    void load(Binding binding, Map<Class<?>, Object> contextObjects);
 
     void unload();