Scripting: implement support for user-defined execution wrapper

libraries/scripting/common/src/kotlin/script/experimental/api/scriptEvaluation.kt

@@ -96,6 +96,24 @@ val ScriptEvaluationConfigurationKeys.hostConfiguration by PropertiesCollection.
  */
 val ScriptEvaluationConfigurationKeys.refineConfigurationBeforeEvaluate by PropertiesCollection.key<List<RefineEvaluationConfigurationData>>(isTransient = true)
 
+interface ScriptExecutionWrapper<T> {
+    fun invoke(block: () -> T): T
+}
+
+/**
+ *  An optional user-defined wrapper which is called with the code that actually executes script body
+ */
+val ScriptEvaluationConfigurationKeys.scriptExecutionWrapper by PropertiesCollection.key<ScriptExecutionWrapper<*>>(isTransient = true)
+
+/**
+ * A helper to enable passing lambda directly to the scriptExecutionWrapper "keyword"
+ */
+fun <T> ScriptEvaluationConfiguration.Builder.scriptExecutionWrapper(wrapper: (() -> T) -> T) {
+    ScriptEvaluationConfiguration.scriptExecutionWrapper.put(object : ScriptExecutionWrapper<T> {
+        override fun invoke(block: () -> T): T = wrapper(block)
+    })
+}
+
 /**
  * A helper to enable scriptsInstancesSharingMap with default implementation
  */

libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt

@@ -290,6 +290,31 @@ class ScriptingHostTest : TestCase() {
         Assert.assertEquals(greeting, output)
     }
 
+    @Test
+    fun testEvalWithWrapper() {
+        val greeting = "Hello from script!"
+        var output = ""
+        BasicJvmScriptingHost().evalWithTemplate<SimpleScriptTemplate>(
+            "println(\"$greeting\")".toScriptSource(),
+            {},
+            {
+                scriptExecutionWrapper {
+                    val outStream = ByteArrayOutputStream()
+                    val prevOut = System.out
+                    System.setOut(PrintStream(outStream))
+                    try {
+                        it()
+                    } finally {
+                        System.out.flush()
+                        System.setOut(prevOut)
+                        output = outStream.toString().trim()
+                    }
+                }
+            }
+        ).throwOnFailure()
+        Assert.assertEquals(greeting, output)
+    }
+
     private fun doDiamondImportTest(evaluationConfiguration: ScriptEvaluationConfiguration? = null): List<String> {
         val mainScript = "sharedVar += 1\nprintln(\"sharedVar == \$sharedVar\")".toScriptSource("main.kts")
         val middleScript = File(TEST_DATA_DIR, "importTest/diamondImportMiddle.kts").toScriptSource()
@@ -582,4 +607,3 @@ internal fun captureOutAndErr(body: () -> Unit): Pair<String, String> {
     }
     return outStream.toString().trim() to errStream.toString().trim()
 }
-

libraries/scripting/jvm/src/kotlin/script/experimental/jvm/BasicJvmScriptEvaluator.kt

@@ -94,10 +94,18 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
 
         val ctor = java.constructors.single()
 
+        @Suppress("UNCHECKED_CAST")
+        val wrapper: ScriptExecutionWrapper<Any>? =
+            refinedEvalConfiguration[ScriptEvaluationConfiguration.scriptExecutionWrapper] as ScriptExecutionWrapper<Any>?
+
         val saveClassLoader = Thread.currentThread().contextClassLoader
         Thread.currentThread().contextClassLoader = this.java.classLoader
         return try {
-            ctor.newInstance(*args.toArray())
+            if (wrapper == null) {
+                ctor.newInstance(*args.toArray())
+            } else wrapper.invoke {
+                ctor.newInstance(*args.toArray())
+            }
         } finally {
             Thread.currentThread().contextClassLoader = saveClassLoader
         }