feat(python): streaming bin script output (#3794)

Resolves #3782.

Inverts control on bin script execution by returning the script to execute (as well as arguments and ENV variables), rather than running it within the kernel process. This allows the calling process to pipe the output streams, rather than the kernel buffering the output and returning it at the end of the synchronous call.

Also, adds more complete tests to the python runtime to validate `stderr`, exit codes and streaming response.

---

By submitting this pull request, I confirm that my contribution is made under the terms of the [Apache 2.0 license].

[Apache 2.0 license]: https://www.apache.org/licenses/LICENSE-2.0

packages/@jsii/kernel/src/api.ts

@@ -97,6 +97,18 @@ export interface LoadResponse {
   readonly types: number;
 }
 
+export interface GetScriptCommandRequest {
+  readonly assembly: string;
+  readonly script: string;
+  readonly args?: string[];
+}
+
+export interface GetScriptCommandResponse {
+  command: string;
+  args: string[];
+  env: Record<string, string>;
+}
+
 export interface InvokeScriptRequest {
   readonly assembly: string;
   readonly script: string;

packages/@jsii/kernel/src/kernel.test.ts

@@ -78,6 +78,8 @@ defineTest.skip = function (
   return defineTest(name, method, test.skip);
 };
 
+defineTest.skipIf = (expr: boolean) => (expr ? defineTest.skip : defineTest);
+
 // Note: this test asserts file permissions, which work differently on Windows, so we skip it there
 (process.platform === 'win32' ? test.skip : test)(
   'load preserves file permissions',
@@ -2138,6 +2140,49 @@ defineTest('Override transitive property', (sandbox) => {
   expect(propValue).toBe('N3W');
 });
 
+defineTest.skipIf(!!recordingOutput)(
+  'getBinScriptCommand() returns output',
+  (sandbox) => {
+    const result = sandbox.getBinScriptCommand({
+      assembly: 'jsii-calc',
+      script: 'calc',
+    });
+
+    expect(result).toMatchObject<api.GetScriptCommandResponse>({
+      command: expect.stringMatching(
+        /node_modules[/\\]jsii-calc[/\\]bin[/\\]run$/,
+      ),
+      args: [],
+      env: {
+        NODE_OPTIONS: expect.any(String),
+        PATH: expect.any(String),
+      },
+    });
+  },
+);
+
+defineTest.skipIf(!!recordingOutput)(
+  'getBinScriptCommand() accepts arguments',
+  (sandbox) => {
+    const result = sandbox.getBinScriptCommand({
+      assembly: 'jsii-calc',
+      script: 'calc',
+      args: ['arg1', 'arg2'],
+    });
+
+    expect(result).toMatchObject<api.GetScriptCommandResponse>({
+      command: expect.stringMatching(
+        /node_modules[/\\]jsii-calc[/\\]bin[/\\]run$/,
+      ),
+      args: ['arg1', 'arg2'],
+      env: {
+        NODE_OPTIONS: expect.any(String),
+        PATH: expect.any(String),
+      },
+    });
+  },
+);
+
 defineTest('invokeBinScript() return output', (sandbox) => {
   const result = sandbox.invokeBinScript({
     assembly: 'jsii-calc',

packages/@jsii/kernel/src/kernel.ts

@@ -169,50 +169,29 @@ export class Kernel {
     };
   }
 
+  public getBinScriptCommand(
+    req: api.GetScriptCommandRequest,
+  ): api.GetScriptCommandResponse {
+    return this._getBinScriptCommand(req);
+  }
+
   public invokeBinScript(
     req: api.InvokeScriptRequest,
   ): api.InvokeScriptResponse {
-    const packageDir = this._getPackageDir(req.assembly);
-    if (fs.pathExistsSync(packageDir)) {
-      // module exists, verify version
-      const epkg = fs.readJsonSync(path.join(packageDir, 'package.json'));
-
-      if (!epkg.bin) {
-        throw new JsiiFault(
-          'There is no bin scripts defined for this package.',
-        );
-      }
-
-      const scriptPath = epkg.bin[req.script];
+    const { command, args, env } = this._getBinScriptCommand(req);
 
-      if (!epkg.bin) {
-        throw new JsiiFault(`Script with name ${req.script} was not defined.`);
-      }
-
-      const result = cp.spawnSync(
-        path.join(packageDir, scriptPath),
-        req.args ?? [],
-        {
-          encoding: 'utf-8',
-          env: {
-            ...process.env,
-            // Make sure the current NODE_OPTIONS are honored if we shell out to node
-            NODE_OPTIONS: process.execArgv.join(' '),
-            // Make sure "this" node is ahead of $PATH just in case
-            PATH: `${path.dirname(process.execPath)}:${process.env.PATH}`,
-          },
-          shell: true,
-        },
-      );
+    const result = cp.spawnSync(command, args, {
+      encoding: 'utf-8',
+      env,
+      shell: true,
+    });
 
-      return {
-        stdout: result.stdout,
-        stderr: result.stderr,
-        status: result.status,
-        signal: result.signal,
-      };
-    }
-    throw new JsiiFault(`Package with name ${req.assembly} was not loaded.`);
+    return {
+      stdout: result.stdout,
+      stderr: result.stderr,
+      status: result.status,
+      signal: result.signal,
+    };
   }
 
   public create(req: api.CreateRequest): api.CreateResponse {
@@ -1325,6 +1304,38 @@ export class Kernel {
     return property;
   }
 
+  /**
+   * Shared (non-public implementation) to as not to break API recording.
+   */
+  private _getBinScriptCommand(
+    req: api.GetScriptCommandRequest,
+  ): api.GetScriptCommandResponse {
+    const packageDir = this._getPackageDir(req.assembly);
+    if (fs.pathExistsSync(packageDir)) {
+      // module exists, verify version
+      const epkg = fs.readJsonSync(path.join(packageDir, 'package.json'));
+
+      const scriptPath = epkg.bin?.[req.script];
+
+      if (!epkg.bin) {
+        throw new JsiiFault(`Script with name ${req.script} was not defined.`);
+      }
+
+      return {
+        command: path.join(packageDir, scriptPath),
+        args: req.args ?? [],
+        env: {
+          ...process.env,
+          // Make sure the current NODE_OPTIONS are honored if we shell out to node
+          NODE_OPTIONS: process.execArgv.join(' '),
+          // Make sure "this" node is ahead of $PATH just in case
+          PATH: `${path.dirname(process.execPath)}:${process.env.PATH}`,
+        },
+      };
+    }
+    throw new JsiiFault(`Package with name ${req.assembly} was not loaded.`);
+  }
+
   //
   // type information
   //

packages/@jsii/python-runtime/src/jsii/_kernel/__init__.py

@@ -14,7 +14,6 @@
 from .. import _reference_map
 from .._utils import Singleton
 from .providers import BaseProvider, ProcessProvider
-from .types import Callback
 from .types import (
     EnumRef,
     LoadRequest,
@@ -34,6 +33,8 @@
     GetResponse,
     InvokeRequest,
     InvokeResponse,
+    GetScriptCommandRequest,
+    GetScriptCommandResponse,
     InvokeScriptRequest,
     InvokeScriptResponse,
     KernelResponse,
@@ -299,6 +300,20 @@ def __init__(self, provider_class: Type[BaseProvider] = ProcessProvider) -> None
     def load(self, name: str, version: str, tarball: str) -> None:
         self.provider.load(LoadRequest(name=name, version=version, tarball=tarball))
 
+    def getBinScriptCommand(
+        self, pkgname: str, script: str, args: Optional[Sequence[str]] = None
+    ) -> GetScriptCommandResponse:
+        if args is None:
+            args = []
+
+        return self.provider.getScriptCommand(
+            GetScriptCommandRequest(
+                assembly=pkgname,
+                script=script,
+                args=_make_reference_for_native(self, args),
+            )
+        )
+
     def invokeBinScript(
         self, pkgname: str, script: str, args: Optional[Sequence[str]] = None
     ) -> InvokeScriptResponse:

packages/@jsii/python-runtime/src/jsii/_kernel/providers/base.py

@@ -11,6 +11,8 @@
     GetResponse,
     InvokeRequest,
     InvokeResponse,
+    GetScriptCommandRequest,
+    GetScriptCommandResponse,
     InvokeScriptRequest,
     InvokeScriptResponse,
     DeleteRequest,
@@ -47,6 +49,12 @@ class BaseProvider(metaclass=abc.ABCMeta):
     def load(self, request: LoadRequest) -> LoadResponse:
         ...
 
+    @abc.abstractmethod
+    def getScriptCommand(
+        self, request: GetScriptCommandRequest
+    ) -> GetScriptCommandResponse:
+        ...
+
     @abc.abstractmethod
     def invokeBinScript(self, request: InvokeScriptRequest) -> InvokeScriptResponse:
         ...

packages/@jsii/python-runtime/src/jsii/_kernel/providers/process.py

@@ -41,6 +41,8 @@
     GetResponse,
     InvokeRequest,
     InvokeResponse,
+    GetScriptCommandRequest,
+    GetScriptCommandResponse,
     InvokeScriptRequest,
     InvokeScriptResponse,
     SetRequest,
@@ -160,6 +162,12 @@ def __init__(self):
             LoadRequest,
             _with_api_key("load", self._serializer.unstructure_attrs_asdict),
         )
+        self._serializer.register_unstructure_hook(
+            GetScriptCommandRequest,
+            _with_api_key(
+                "getBinScriptCommand", self._serializer.unstructure_attrs_asdict
+            ),
+        )
         self._serializer.register_unstructure_hook(
             InvokeScriptRequest,
             _with_api_key("invokeBinScript", self._serializer.unstructure_attrs_asdict),
@@ -343,6 +351,11 @@ def _process(self) -> _NodeProcess:
     def load(self, request: LoadRequest) -> LoadResponse:
         return self._process.send(request, LoadResponse)
 
+    def getScriptCommand(
+        self, request: GetScriptCommandRequest
+    ) -> GetScriptCommandResponse:
+        return self._process.send(request, GetScriptCommandResponse)
+
     def invokeBinScript(self, request: InvokeScriptRequest) -> InvokeScriptResponse:
         return self._process.send(request, InvokeScriptResponse)
 

packages/@jsii/python-runtime/src/jsii/_kernel/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Generic, List, Optional, Mapping, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Mapping, TypeVar, Union
 from typing_extensions import Protocol
 
 import attr
@@ -41,6 +41,22 @@ class LoadResponse:
     types: int
 
 
+@attr.s(auto_attribs=True, frozen=True, slots=True)
+class GetScriptCommandRequest:
+
+    assembly: str
+    script: str
+    args: List[Any] = attr.Factory(list)
+
+
+@attr.s(auto_attribs=True, frozen=True, slots=True)
+class GetScriptCommandResponse:
+
+    command: str
+    args: List[str] = attr.Factory(list)
+    env: Dict[str, str] = attr.Factory(dict)
+
+
 @attr.s(auto_attribs=True, frozen=True, slots=True)
 class InvokeScriptRequest:
 

packages/@jsii/python-runtime/src/jsii/_runtime.py

@@ -1,6 +1,7 @@
 import abc
 import os
 import sys
+import subprocess
 
 import attr
 
@@ -66,14 +67,15 @@ def invokeBinScript(
         if args is None:
             args = []
 
-        response = _kernel.invokeBinScript(pkgname, script, args)
-        if response.stdout:
-            print(response.stdout)
+        response = _kernel.getBinScriptCommand(pkgname, script, args)
 
-        if response.stderr:
-            print(response.stderr, file=sys.stderr)
-
-        return response.status
+        result = subprocess.run(
+            " ".join([response.command, *response.args]),
+            encoding="utf-8",
+            shell=True,
+            env=response.env,
+        )
+        return result.returncode
 
 
 class JSIIMeta(_ClassPropertyMeta, type):