Change GroovyScriptExtension.load to use a map of context objects

afs-action-dsl/src/main/java/com/powsybl/afs/action/dsl/ActionScript.java

@@ -17,6 +17,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.ResourceBundle;
 
@@ -51,8 +52,18 @@ public ActionDb load(Network network) {
         return new ActionDslLoader(readScript(true)).load(network);
     }
 
+    public ActionDb load(Network network, Map<Class<?>, Object> contextObjects) {
+        Objects.requireNonNull(network);
+        return new ActionDslLoader(readScript(true)).load(network, contextObjects);
+    }
+
     @Override
     public List<Contingency> getContingencies(Network network) {
         return new ArrayList<>(load(network).getContingencies());
     }
+
+    @Override
+    public List<Contingency> getContingencies(Network network, Map<Class<?>, Object> contextObjects) {
+        return new ArrayList<>(load(network, contextObjects).getContingencies());
+    }
 }

afs-ext-base/pom.xml

@@ -115,6 +115,10 @@
             <artifactId>powsybl-iidm-impl</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.powsybl</groupId>
+            <artifactId>powsybl-scripting</artifactId>
+        </dependency>
     </dependencies>
 
 </project>

afs-ext-base/src/main/java/com/powsybl/afs/ext/base/AbstractProjectCase.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2025, RTE (https://www.rte-france.com)
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ * SPDX-License-Identifier: MPL-2.0
+ */
+package com.powsybl.afs.ext.base;
+
+import com.powsybl.afs.ProjectFile;
+import com.powsybl.afs.ProjectFileCreationContext;
+import com.powsybl.iidm.network.Network;
+import com.powsybl.iidm.network.NetworkListener;
+import com.powsybl.scripting.groovy.GroovyScriptExtension;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * @author Nicolas Rol {@literal <nicolas.rol at rte-france.com>}
+ */
+public abstract class AbstractProjectCase extends ProjectFile implements ProjectCase {
+
+    protected AbstractProjectCase(ProjectFileCreationContext context, int version) {
+        super(context, version);
+    }
+
+    @Override
+    public String queryNetwork(ScriptType scriptType, String scriptContent) {
+        return queryNetwork(scriptType, scriptContent, Collections.emptyList(), Collections.emptyMap());
+    }
+
+    @Override
+    public String queryNetwork(ScriptType scriptType, String scriptContent, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        Objects.requireNonNull(scriptType);
+        Objects.requireNonNull(scriptContent);
+        return findService(NetworkCacheService.class).queryNetwork(this, scriptType, scriptContent, extensions, contextObjects);
+    }
+
+    @Override
+    public Network getNetwork() {
+        return findService(NetworkCacheService.class).getNetwork(this);
+    }
+
+    @Override
+    public Network getNetwork(Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        return findService(NetworkCacheService.class).getNetwork(this, extensions, contextObjects);
+    }
+
+    @Override
+    public Network getNetwork(List<NetworkListener> listeners) {
+        return findService(NetworkCacheService.class).getNetwork(this, listeners);
+    }
+
+    @Override
+    public Network getNetwork(Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects, List<NetworkListener> listeners) {
+        return findService(NetworkCacheService.class).getNetwork(this, listeners, extensions, contextObjects);
+    }
+
+    @Override
+    public void invalidateNetworkCache() {
+        findService(NetworkCacheService.class).invalidateCache(this);
+    }
+}

afs-ext-base/src/main/java/com/powsybl/afs/ext/base/ImportedCase.java

@@ -14,7 +14,6 @@
 import com.powsybl.iidm.network.Importer;
 import com.powsybl.iidm.network.ImportersLoader;
 import com.powsybl.iidm.network.Network;
-import com.powsybl.iidm.network.NetworkListener;
 
 import java.io.IOException;
 import java.io.InputStreamReader;
@@ -31,7 +30,7 @@
  *
  *  @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
  */
-public class ImportedCase extends ProjectFile implements ProjectCase {
+public class ImportedCase extends AbstractProjectCase {
 
     public static final String PSEUDO_CLASS = "importedCase";
     public static final int VERSION = 0;
@@ -69,28 +68,6 @@ public Importer getImporter() {
                 .orElseThrow(() -> new AfsException("Importer not found for format " + format));
     }
 
-    @Override
-    public String queryNetwork(ScriptType scriptType, String scriptContent) {
-        Objects.requireNonNull(scriptType);
-        Objects.requireNonNull(scriptContent);
-        return findService(NetworkCacheService.class).queryNetwork(this, scriptType, scriptContent);
-    }
-
-    @Override
-    public Network getNetwork() {
-        return findService(NetworkCacheService.class).getNetwork(this);
-    }
-
-    @Override
-    public Network getNetwork(List<NetworkListener> listeners) {
-        return findService(NetworkCacheService.class).getNetwork(this, listeners);
-    }
-
-    @Override
-    public void invalidateNetworkCache() {
-        findService(NetworkCacheService.class).invalidateCache(this);
-    }
-
     @Override
     protected List<ProjectFile> invalidate() {
         List<ProjectFile> dependencies = super.invalidate();

afs-ext-base/src/main/java/com/powsybl/afs/ext/base/LocalNetworkCacheService.java

@@ -13,11 +13,17 @@
 import com.powsybl.iidm.network.Network;
 import com.powsybl.iidm.network.NetworkFactory;
 import com.powsybl.iidm.network.NetworkListener;
+import com.powsybl.scripting.groovy.GroovyScriptExtension;
 import groovy.json.JsonOutput;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.*;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Properties;
+import java.util.UUID;
 
 /**
  * @author Geoffroy Jamgotchian {@literal <geoffroy.jamgotchian at rte-france.com>}
@@ -29,11 +35,11 @@ public class LocalNetworkCacheService implements NetworkCacheService {
     private final ScriptCache<ProjectFile, Network, ProjectCaseListener> cache;
 
     public LocalNetworkCacheService() {
-        cache = new ScriptCache<>(50, 1, projectFile -> {
+        cache = new ScriptCache<>(50, 1, (projectFile, extensions, contextObjects) -> {
             UUID taskId = projectFile.startTask();
             try {
                 projectFile.createLogger(taskId).log("Loading network...");
-                return loadNetworkFromProjectCase((ProjectCase) projectFile);
+                return loadNetworkFromProjectCase((ProjectCase) projectFile, extensions, contextObjects);
             } finally {
                 projectFile.stopTask(taskId);
             }
@@ -55,8 +61,8 @@ private static ScriptResult<Network> loadNetworkFromImportedCase(ImportedCase im
         return ScriptResult.of(network);
     }
 
-    private static ScriptResult<Network> applyScript(Network network, String previousScriptOutput, ModificationScript script) {
-        ScriptResult<Object> result = ScriptUtils.runScript(network, script.getScriptType(), script.readScript(true));
+    private static ScriptResult<Network> applyScript(Network network, String previousScriptOutput, ModificationScript script, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        ScriptResult<Object> result = ScriptUtils.runScript(network, script.getScriptType(), script.readScript(true), extensions, contextObjects);
         if (result.getError() == null) {
             return new ScriptResult<>(network, previousScriptOutput + result.getOutput(), null);
         } else {
@@ -65,47 +71,53 @@ private static ScriptResult<Network> applyScript(Network network, String previou
         }
     }
 
-    private static ScriptResult<Network> loadNetworkFromVirtualCase(VirtualCase virtualCase, List<NetworkListener> listeners) {
-        ProjectCase baseCase = (ProjectCase) virtualCase.getCase()
-                                                        .orElseThrow(() -> new AfsException("Case link is dead"));
+    private static ScriptResult<Network> loadNetworkFromVirtualCase(VirtualCase virtualCase, List<NetworkListener> listeners, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        ProjectCase baseCase = (ProjectCase) virtualCase.getCase().orElseThrow(() -> new AfsException("Case link is dead"));
 
         ScriptResult<Network> network = loadNetworkFromProjectCase(baseCase, listeners);
 
         if (network.getError() != null) {
             return network;
         }
 
-        ModificationScript script = virtualCase.getScript()
-                                               .orElseThrow(VirtualCase::createScriptLinkIsDeadException);
+        ModificationScript script = virtualCase.getScript().orElseThrow(VirtualCase::createScriptLinkIsDeadException);
 
         LOGGER.info("Applying script to network of project case {}", virtualCase.getId());
 
-        return applyScript(network.getValue(), network.getOutput(), script);
+        return applyScript(network.getValue(), network.getOutput(), script, extensions, contextObjects);
     }
 
-    private static ScriptResult<Network> loadNetworkFromProjectCase(ProjectCase projectCase) {
-        return loadNetworkFromProjectCase(projectCase, Collections.emptyList());
+    private static ScriptResult<Network> loadNetworkFromProjectCase(ProjectCase projectCase, List<NetworkListener> listeners) {
+        return loadNetworkFromProjectCase(projectCase, listeners, Collections.emptyList(), Collections.emptyMap());
     }
 
-    private static ScriptResult<Network> loadNetworkFromProjectCase(ProjectCase projectCase, List<NetworkListener> listeners) {
+    private static ScriptResult<Network> loadNetworkFromProjectCase(ProjectCase projectCase, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        return loadNetworkFromProjectCase(projectCase, Collections.emptyList(), extensions, contextObjects);
+    }
+
+    private static ScriptResult<Network> loadNetworkFromProjectCase(ProjectCase projectCase, List<NetworkListener> listeners, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
         if (projectCase instanceof ImportedCase importedCase) {
             return loadNetworkFromImportedCase(importedCase, listeners);
         } else if (projectCase instanceof VirtualCase virtualCase) {
-            return loadNetworkFromVirtualCase(virtualCase, listeners);
+            return loadNetworkFromVirtualCase(virtualCase, listeners, extensions, contextObjects);
         } else {
             throw new AssertionError("ProjectCase implementation " + projectCase.getClass().getName() + " not supported");
         }
     }
 
     @Override
     public <T extends ProjectFile & ProjectCase> String queryNetwork(T projectCase, ScriptType scriptType, String scriptContent) {
+        return queryNetwork(projectCase, scriptType, scriptContent, Collections.emptyList(), Collections.emptyMap());
+    }
+
+    @Override
+    public <T extends ProjectFile & ProjectCase> String queryNetwork(T projectCase, ScriptType scriptType, String scriptContent, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
         Objects.requireNonNull(projectCase);
         Objects.requireNonNull(scriptType);
         Objects.requireNonNull(scriptContent);
 
         Network network = getNetwork(projectCase);
-
-        ScriptResult<Object> result = ScriptUtils.runScript(network, ScriptType.GROOVY, scriptContent);
+        ScriptResult<Object> result = ScriptUtils.runScript(network, ScriptType.GROOVY, scriptContent, extensions, contextObjects);
         if (result.getError() != null) {
             throw new ScriptException(projectCase, result.getError());
         }
@@ -114,12 +126,22 @@ public <T extends ProjectFile & ProjectCase> String queryNetwork(T projectCase,
 
     @Override
     public <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase) {
-        return cache.get(projectCase).getValueOrThrowIfError(projectCase);
+        return getNetwork(projectCase, Collections.emptyList(), Collections.emptyMap());
+    }
+
+    @Override
+    public <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        return cache.get(projectCase, extensions, contextObjects).getValueOrThrowIfError(projectCase);
     }
 
     @Override
     public <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, List<NetworkListener> listeners) {
-        ScriptResult<Network> network = loadNetworkFromProjectCase(projectCase, listeners);
+        return getNetwork(projectCase, listeners, Collections.emptyList(), Collections.emptyMap());
+    }
+
+    @Override
+    public <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, List<NetworkListener> listeners, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects) {
+        ScriptResult<Network> network = loadNetworkFromProjectCase(projectCase, listeners, extensions, contextObjects);
         return network.getValueOrThrowIfError(projectCase);
     }
 

afs-ext-base/src/main/java/com/powsybl/afs/ext/base/NetworkCacheService.java

@@ -9,9 +9,11 @@
 import com.powsybl.afs.ProjectFile;
 import com.powsybl.iidm.network.Network;
 import com.powsybl.iidm.network.NetworkListener;
+import com.powsybl.scripting.groovy.GroovyScriptExtension;
 import org.apache.commons.lang3.StringUtils;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * Provides caching capabilities for loaded {@code Network} objects.
@@ -22,10 +24,16 @@ public interface NetworkCacheService {
 
     <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase);
 
+    <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects);
+
     <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, List<NetworkListener> listeners);
 
+    <T extends ProjectFile & ProjectCase> Network getNetwork(T projectCase, List<NetworkListener> listeners, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects);
+
     <T extends ProjectFile & ProjectCase> String queryNetwork(T projectCase, ScriptType scriptType, String scriptContent);
 
+    <T extends ProjectFile & ProjectCase> String queryNetwork(T projectCase, ScriptType scriptType, String scriptContent, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects);
+
     <T extends ProjectFile & ProjectCase> void invalidateCache(T projectCase);
 
     <T extends ProjectFile & ProjectCase> void addListener(T projectCase, ProjectCaseListener listener);

afs-ext-base/src/main/java/com/powsybl/afs/ext/base/ProjectCase.java

@@ -8,8 +8,10 @@
 
 import com.powsybl.iidm.network.Network;
 import com.powsybl.iidm.network.NetworkListener;
+import com.powsybl.scripting.groovy.GroovyScriptExtension;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * Common interface for project files able to provide a Network.
@@ -20,8 +22,12 @@ public interface ProjectCase {
 
     String queryNetwork(ScriptType scriptType, String scriptContent);
 
+    String queryNetwork(ScriptType scriptType, String scriptContent, Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects);
+
     Network getNetwork();
 
+    Network getNetwork(Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects);
+
     /**
      * Get the network and add a listeners on it in order to listen changes due to virtual case script application.
      * The listeners will not be removed from the network at the end of the network loading,
@@ -31,6 +37,8 @@ public interface ProjectCase {
      */
     Network getNetwork(List<NetworkListener> listeners);
 
+    Network getNetwork(Iterable<GroovyScriptExtension> extensions, Map<Class<?>, Object> contextObjects, List<NetworkListener> listeners);
+
     void invalidateNetworkCache();
 
     void addListener(ProjectCaseListener l);