From 6784cb63aa42294a4dc860eddb93d8e1adb44eb3 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Fri, 13 May 2022 17:16:43 +0200
Subject: [PATCH] Scripting: implement support for user-defined execution
 wrapper

---
 .../experimental/api/scriptEvaluation.kt      | 18 +++++++++++++
 .../jvmhost/test/ScriptingHostTest.kt         | 26 ++++++++++++++++++-
 .../jvm/BasicJvmScriptEvaluator.kt            | 10 ++++++-
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/libraries/scripting/common/src/kotlin/script/experimental/api/scriptEvaluation.kt b/libraries/scripting/common/src/kotlin/script/experimental/api/scriptEvaluation.kt
index c9e0f826679..ac39d2326d0 100644
--- a/libraries/scripting/common/src/kotlin/script/experimental/api/scriptEvaluation.kt
+++ b/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
  */
diff --git a/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt b/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt
index f46d5e7dedb..1a877eb0450 100644
--- a/libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ScriptingHostTest.kt
+++ b/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()
 }
-
diff --git a/libraries/scripting/jvm/src/kotlin/script/experimental/jvm/BasicJvmScriptEvaluator.kt b/libraries/scripting/jvm/src/kotlin/script/experimental/jvm/BasicJvmScriptEvaluator.kt
index 49de86af32f..2ad4018a1d3 100644
--- a/libraries/scripting/jvm/src/kotlin/script/experimental/jvm/BasicJvmScriptEvaluator.kt
+++ b/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
         }