Use stdin if workspace has large number of requirements (#21988)

Closes #21480

Author: karthiknadig

pythonFiles/create_venv.py

@@ -3,6 +3,7 @@
 
 import argparse
 import importlib.util as import_util
+import json
 import os
 import pathlib
 import subprocess
@@ -56,6 +57,12 @@ def parse_args(argv: Sequence[str]) -> argparse.Namespace:
         metavar="NAME",
         action="store",
     )
+    parser.add_argument(
+        "--stdin",
+        action="store_true",
+        default=False,
+        help="Read arguments from stdin.",
+    )
     return parser.parse_args(argv)
 
 
@@ -152,6 +159,16 @@ def install_pip(name: str):
     )
 
 
+def get_requirements_from_args(args: argparse.Namespace) -> List[str]:
+    requirements = []
+    if args.stdin:
+        data = json.loads(sys.stdin.read())
+        requirements = data.get("requirements", [])
+    if args.requirements:
+        requirements.extend(args.requirements)
+    return requirements
+
+
 def main(argv: Optional[Sequence[str]] = None) -> None:
     if argv is None:
         argv = []
@@ -223,9 +240,10 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
         print(f"VENV_INSTALLING_PYPROJECT: {args.toml}")
         install_toml(venv_path, args.extras)
 
-    if args.requirements:
-        print(f"VENV_INSTALLING_REQUIREMENTS: {args.requirements}")
-        install_requirements(venv_path, args.requirements)
+    requirements = get_requirements_from_args(args)
+    if requirements:
+        print(f"VENV_INSTALLING_REQUIREMENTS: {requirements}")
+        install_requirements(venv_path, requirements)
 
 
 if __name__ == "__main__":

pythonFiles/tests/test_create_venv.py

@@ -1,7 +1,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+import argparse
+import contextlib
 import importlib
+import io
+import json
 import os
 import sys
 
@@ -224,3 +228,50 @@ def run_process(args, error_message):
 
     create_venv.run_process = run_process
     create_venv.main([])
+
+
+@contextlib.contextmanager
+def redirect_io(stream: str, new_stream):
+    """Redirect stdio streams to a custom stream."""
+    old_stream = getattr(sys, stream)
+    setattr(sys, stream, new_stream)
+    yield
+    setattr(sys, stream, old_stream)
+
+
+class CustomIO(io.TextIOWrapper):
+    """Custom stream object to replace stdio."""
+
+    name: str = "customio"
+
+    def __init__(self, name: str, encoding="utf-8", newline=None):
+        self._buffer = io.BytesIO()
+        self._buffer.name = name
+        super().__init__(self._buffer, encoding=encoding, newline=newline)
+
+    def close(self):
+        """Provide this close method which is used by some tools."""
+        # This is intentionally empty.
+
+    def get_value(self) -> str:
+        """Returns value from the buffer as string."""
+        self.seek(0)
+        return self.read()
+
+
+def test_requirements_from_stdin():
+    importlib.reload(create_venv)
+
+    cli_requirements = [f"cli-requirement{i}.txt" for i in range(3)]
+    args = argparse.Namespace()
+    args.__dict__.update({"stdin": True, "requirements": cli_requirements})
+
+    stdin_requirements = [f"stdin-requirement{i}.txt" for i in range(20)]
+    text = json.dumps({"requirements": stdin_requirements})
+    str_input = CustomIO("<stdin>", encoding="utf-8", newline="\n")
+    with redirect_io("stdin", str_input):
+        str_input.write(text)
+        str_input.seek(0)
+        actual = create_venv.get_requirements_from_args(args)
+
+    assert actual == stdin_requirements + cli_requirements

src/client/pythonEnvironments/creation/provider/venvCreationProvider.ts

@@ -32,8 +32,14 @@ import {
     CreateEnvironmentResult,
 } from '../proposed.createEnvApis';
 
-function generateCommandArgs(installInfo?: IPackageInstallSelection[], addGitIgnore?: boolean): string[] {
+interface IVenvCommandArgs {
+    argv: string[];
+    stdin: string | undefined;
+}
+
+function generateCommandArgs(installInfo?: IPackageInstallSelection[], addGitIgnore?: boolean): IVenvCommandArgs {
     const command: string[] = [createVenvScript()];
+    let stdin: string | undefined;
 
     if (addGitIgnore) {
         command.push('--git-ignore');
@@ -52,14 +58,21 @@ function generateCommandArgs(installInfo?: IPackageInstallSelection[], addGitIgn
         });
 
         const requirements = installInfo.filter((i) => i.installType === 'requirements').map((i) => i.installItem);
-        requirements.forEach((r) => {
-            if (r) {
-                command.push('--requirements', r);
-            }
-        });
+
+        if (requirements.length < 10) {
+            requirements.forEach((r) => {
+                if (r) {
+                    command.push('--requirements', r);
+                }
+            });
+        } else {
+            command.push('--stdin');
+            // Too many requirements can cause the command line to be too long error.
+            stdin = JSON.stringify({ requirements });
+        }
     }
 
-    return command;
+    return { argv: command, stdin };
 }
 
 function getVenvFromOutput(output: string): string | undefined {
@@ -81,7 +94,7 @@ function getVenvFromOutput(output: string): string | undefined {
 async function createVenv(
     workspace: WorkspaceFolder,
     command: string,
-    args: string[],
+    args: IVenvCommandArgs,
     progress: CreateEnvironmentProgress,
     token?: CancellationToken,
 ): Promise<string | undefined> {
@@ -94,11 +107,15 @@ async function createVenv(
     });
 
     const deferred = createDeferred<string | undefined>();
-    traceLog('Running Env creation script: ', [command, ...args]);
-    const { proc, out, dispose } = execObservable(command, args, {
+    traceLog('Running Env creation script: ', [command, ...args.argv]);
+    if (args.stdin) {
+        traceLog('Requirements passed in via stdin: ', args.stdin);
+    }
+    const { proc, out, dispose } = execObservable(command, args.argv, {
         mergeStdOutErr: true,
         token,
         cwd: workspace.uri.fsPath,
+        stdinStr: args.stdin,
     });
 
     const progressAndTelemetry = new VenvProgressAndTelemetry(progress);

src/test/pythonEnvironments/creation/provider/venvCreationProvider.unit.test.ts

@@ -15,7 +15,7 @@ import * as rawProcessApis from '../../../../client/common/process/rawProcessApi
 import * as commonUtils from '../../../../client/pythonEnvironments/creation/common/commonUtils';
 import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../../../constants';
 import { createDeferred } from '../../../../client/common/utils/async';
-import { Output } from '../../../../client/common/process/types';
+import { Output, SpawnOptions } from '../../../../client/common/process/types';
 import { VENV_CREATED_MARKER } from '../../../../client/pythonEnvironments/creation/provider/venvProgressAndTelemetry';
 import { CreateEnv } from '../../../../client/common/utils/localize';
 import * as venvUtils from '../../../../client/pythonEnvironments/creation/provider/venvUtils';
@@ -394,4 +394,157 @@ suite('venv Creation provider tests', () => {
         assert.isTrue(showErrorMessageWithLogsStub.notCalled);
         assert.isTrue(deleteEnvironmentStub.notCalled);
     });
+
+    test('Create venv with 1000 requirement files', async () => {
+        pickWorkspaceFolderStub.resolves(workspace1);
+
+        interpreterQuickPick
+            .setup((i) => i.getInterpreterViaQuickPick(typemoq.It.isAny(), typemoq.It.isAny(), typemoq.It.isAny()))
+            .returns(() => Promise.resolve('/usr/bin/python'))
+            .verifiable(typemoq.Times.once());
+
+        const requirements = Array.from({ length: 1000 }, (_, i) => ({
+            installType: 'requirements',
+            installItem: `requirements${i}.txt`,
+        }));
+        pickPackagesToInstallStub.resolves(requirements);
+        const expected = JSON.stringify({ requirements: requirements.map((r) => r.installItem) });
+
+        const deferred = createDeferred();
+        let _next: undefined | ((value: Output<string>) => void);
+        let _complete: undefined | (() => void);
+        let stdin: undefined | string;
+        let hasStdinArg = false;
+        execObservableStub.callsFake((_c, argv: string[], options) => {
+            stdin = options?.stdinStr;
+            hasStdinArg = argv.includes('--stdin');
+            deferred.resolve();
+            return {
+                proc: {
+                    exitCode: 0,
+                },
+                out: {
+                    subscribe: (
+                        next?: (value: Output<string>) => void,
+                        _error?: (error: unknown) => void,
+                        complete?: () => void,
+                    ) => {
+                        _next = next;
+                        _complete = complete;
+                    },
+                },
+                dispose: () => undefined,
+            };
+        });
+
+        progressMock.setup((p) => p.report({ message: CreateEnv.statusStarting })).verifiable(typemoq.Times.once());
+
+        withProgressStub.callsFake(
+            (
+                _options: ProgressOptions,
+                task: (
+                    progress: CreateEnvironmentProgress,
+                    token?: CancellationToken,
+                ) => Thenable<CreateEnvironmentResult>,
+            ) => task(progressMock.object),
+        );
+
+        const promise = venvProvider.createEnvironment();
+        await deferred.promise;
+        assert.isDefined(_next);
+        assert.isDefined(_complete);
+
+        _next!({ out: `${VENV_CREATED_MARKER}new_environment`, source: 'stdout' });
+        _complete!();
+
+        const actual = await promise;
+        assert.deepStrictEqual(actual, {
+            path: 'new_environment',
+            workspaceFolder: workspace1,
+        });
+        interpreterQuickPick.verifyAll();
+        progressMock.verifyAll();
+        assert.isTrue(showErrorMessageWithLogsStub.notCalled);
+        assert.isTrue(deleteEnvironmentStub.notCalled);
+        assert.strictEqual(stdin, expected);
+        assert.isTrue(hasStdinArg);
+    });
+
+    test('Create venv with 5 requirement files', async () => {
+        pickWorkspaceFolderStub.resolves(workspace1);
+
+        interpreterQuickPick
+            .setup((i) => i.getInterpreterViaQuickPick(typemoq.It.isAny(), typemoq.It.isAny(), typemoq.It.isAny()))
+            .returns(() => Promise.resolve('/usr/bin/python'))
+            .verifiable(typemoq.Times.once());
+
+        const requirements = Array.from({ length: 5 }, (_, i) => ({
+            installType: 'requirements',
+            installItem: `requirements${i}.txt`,
+        }));
+        pickPackagesToInstallStub.resolves(requirements);
+        const expectedRequirements = requirements.map((r) => r.installItem).sort();
+
+        const deferred = createDeferred();
+        let _next: undefined | ((value: Output<string>) => void);
+        let _complete: undefined | (() => void);
+        let stdin: undefined | string;
+        let hasStdinArg = false;
+        let actualRequirements: string[] = [];
+        execObservableStub.callsFake((_c, argv: string[], options: SpawnOptions) => {
+            stdin = options?.stdinStr;
+            actualRequirements = argv.filter((arg) => arg.startsWith('requirements')).sort();
+            hasStdinArg = argv.includes('--stdin');
+            deferred.resolve();
+            return {
+                proc: {
+                    exitCode: 0,
+                },
+                out: {
+                    subscribe: (
+                        next?: (value: Output<string>) => void,
+                        _error?: (error: unknown) => void,
+                        complete?: () => void,
+                    ) => {
+                        _next = next;
+                        _complete = complete;
+                    },
+                },
+                dispose: () => undefined,
+            };
+        });
+
+        progressMock.setup((p) => p.report({ message: CreateEnv.statusStarting })).verifiable(typemoq.Times.once());
+
+        withProgressStub.callsFake(
+            (
+                _options: ProgressOptions,
+                task: (
+                    progress: CreateEnvironmentProgress,
+                    token?: CancellationToken,
+                ) => Thenable<CreateEnvironmentResult>,
+            ) => task(progressMock.object),
+        );
+
+        const promise = venvProvider.createEnvironment();
+        await deferred.promise;
+        assert.isDefined(_next);
+        assert.isDefined(_complete);
+
+        _next!({ out: `${VENV_CREATED_MARKER}new_environment`, source: 'stdout' });
+        _complete!();
+
+        const actual = await promise;
+        assert.deepStrictEqual(actual, {
+            path: 'new_environment',
+            workspaceFolder: workspace1,
+        });
+        interpreterQuickPick.verifyAll();
+        progressMock.verifyAll();
+        assert.isTrue(showErrorMessageWithLogsStub.notCalled);
+        assert.isTrue(deleteEnvironmentStub.notCalled);
+        assert.isUndefined(stdin);
+        assert.deepStrictEqual(actualRequirements, expectedRequirements);
+        assert.isFalse(hasStdinArg);
+    });
 });