Create env shows requirements files or pyproject.toml extras when available (#20524)

Closes #20277
Closes #20278

Author: karthiknadig

package-lock.json

@@ -1806,6 +1806,7 @@
         "webpack": "webpack"
     },
     "dependencies": {
+        "@ltd/j-toml": "^1.37.0",
         "@vscode/extension-telemetry": "^0.7.4-preview",
         "@vscode/jupyter-lsp-middleware": "^0.2.50",
         "arch": "^2.1.0",

package.json

@@ -7,7 +7,7 @@
 import pathlib
 import subprocess
 import sys
-from typing import Optional, Sequence, Union
+from typing import List, Optional, Sequence, Union
 
 VENV_NAME = ".venv"
 CWD = pathlib.PurePath(os.getcwd())
@@ -19,12 +19,27 @@ class VenvError(Exception):
 
 def parse_args(argv: Sequence[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser()
+
     parser.add_argument(
-        "--install",
-        action="store_true",
-        default=False,
-        help="Install packages into the virtual environment.",
+        "--requirements",
+        action="append",
+        default=[],
+        help="Install additional dependencies into the virtual environment.",
+    )
+
+    parser.add_argument(
+        "--toml",
+        action="store",
+        default=None,
+        help="Install additional dependencies from sources like `pyproject.toml` into the virtual environment.",
     )
+    parser.add_argument(
+        "--extras",
+        action="append",
+        default=[],
+        help="Install specific package groups from `pyproject.toml` into the virtual environment.",
+    )
+
     parser.add_argument(
         "--git-ignore",
         action="store_true",
@@ -71,30 +86,36 @@ def get_venv_path(name: str) -> str:
         return os.fspath(CWD / name / "bin" / "python")
 
 
-def install_packages(venv_path: str) -> None:
-    requirements = os.fspath(CWD / "requirements.txt")
-    pyproject = os.fspath(CWD / "pyproject.toml")
+def install_requirements(venv_path: str, requirements: List[str]) -> None:
+    if not requirements:
+        return
 
+    print(f"VENV_INSTALLING_REQUIREMENTS: {requirements}")
+    args = []
+    for requirement in requirements:
+        args += ["-r", requirement]
+    run_process(
+        [venv_path, "-m", "pip", "install"] + args,
+        "CREATE_VENV.PIP_FAILED_INSTALL_REQUIREMENTS",
+    )
+    print("CREATE_VENV.PIP_INSTALLED_REQUIREMENTS")
+
+
+def install_toml(venv_path: str, extras: List[str]) -> None:
+    args = "." if len(extras) == 0 else f".[{','.join(extras)}]"
+    run_process(
+        [venv_path, "-m", "pip", "install", "-e", args],
+        "CREATE_VENV.PIP_FAILED_INSTALL_PYPROJECT",
+    )
+    print("CREATE_VENV.PIP_INSTALLED_PYPROJECT")
+
+
+def upgrade_pip(venv_path: str) -> None:
     run_process(
         [venv_path, "-m", "pip", "install", "--upgrade", "pip"],
         "CREATE_VENV.PIP_UPGRADE_FAILED",
     )
 
-    if file_exists(requirements):
-        print(f"VENV_INSTALLING_REQUIREMENTS: {requirements}")
-        run_process(
-            [venv_path, "-m", "pip", "install", "-r", requirements],
-            "CREATE_VENV.PIP_FAILED_INSTALL_REQUIREMENTS",
-        )
-        print("CREATE_VENV.PIP_INSTALLED_REQUIREMENTS")
-    elif file_exists(pyproject):
-        print(f"VENV_INSTALLING_PYPROJECT: {pyproject}")
-        run_process(
-            [venv_path, "-m", "pip", "install", "-e", ".[extras]"],
-            "CREATE_VENV.PIP_FAILED_INSTALL_PYPROJECT",
-        )
-        print("CREATE_VENV.PIP_INSTALLED_PYPROJECT")
-
 
 def add_gitignore(name: str) -> None:
     git_ignore = CWD / name / ".gitignore"
@@ -112,7 +133,9 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
     if not is_installed("venv"):
         raise VenvError("CREATE_VENV.VENV_NOT_FOUND")
 
-    if args.install and not is_installed("pip"):
+    pip_installed = is_installed("pip")
+    deps_needed = args.requirements or args.extras or args.toml
+    if deps_needed and not pip_installed:
         raise VenvError("CREATE_VENV.PIP_NOT_FOUND")
 
     if venv_exists(args.name):
@@ -128,8 +151,16 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
         if args.git_ignore:
             add_gitignore(args.name)
 
-    if args.install:
-        install_packages(venv_path)
+    if pip_installed:
+        upgrade_pip(venv_path)
+
+    if args.requirements:
+        print(f"VENV_INSTALLING_REQUIREMENTS: {args.requirements}")
+        install_requirements(venv_path, args.requirements)
+
+    if args.toml:
+        print(f"VENV_INSTALLING_PYPROJECT: {args.toml}")
+        install_toml(venv_path, args.extras)
 
 
 if __name__ == "__main__":

pythonFiles/create_venv.py

@@ -16,38 +16,44 @@ def test_venv_not_installed():
     assert str(e.value) == "CREATE_VENV.VENV_NOT_FOUND"
 
 
-def test_pip_not_installed():
+@pytest.mark.parametrize("install", ["requirements", "toml"])
+def test_pip_not_installed(install):
     importlib.reload(create_venv)
     create_venv.venv_exists = lambda _n: True
     create_venv.is_installed = lambda module: module != "pip"
     create_venv.run_process = lambda _args, _error_message: None
     with pytest.raises(create_venv.VenvError) as e:
-        create_venv.main(["--install"])
+        if install == "requirements":
+            create_venv.main(["--requirements", "requirements-for-test.txt"])
+        elif install == "toml":
+            create_venv.main(["--toml", "pyproject.toml", "--extras", "test"])
     assert str(e.value) == "CREATE_VENV.PIP_NOT_FOUND"
 
 
-@pytest.mark.parametrize("env_exists", [True, False])
-@pytest.mark.parametrize("git_ignore", [True, False])
-@pytest.mark.parametrize("install", [True, False])
+@pytest.mark.parametrize("env_exists", ["hasEnv", "noEnv"])
+@pytest.mark.parametrize("git_ignore", ["useGitIgnore", "skipGitIgnore"])
+@pytest.mark.parametrize("install", ["requirements", "toml", "skipInstall"])
 def test_create_env(env_exists, git_ignore, install):
     importlib.reload(create_venv)
     create_venv.is_installed = lambda _x: True
-    create_venv.venv_exists = lambda _n: env_exists
+    create_venv.venv_exists = lambda _n: env_exists == "hasEnv"
+    create_venv.upgrade_pip = lambda _x: None
 
     install_packages_called = False
 
-    def install_packages(_name):
+    def install_packages(_env, _name):
         nonlocal install_packages_called
         install_packages_called = True
 
-    create_venv.install_packages = install_packages
+    create_venv.install_requirements = install_packages
+    create_venv.install_toml = install_packages
 
     run_process_called = False
 
     def run_process(args, error_message):
         nonlocal run_process_called
         run_process_called = True
-        if not env_exists:
+        if env_exists == "noEnv":
             assert args == [sys.executable, "-m", "venv", create_venv.VENV_NAME]
             assert error_message == "CREATE_VENV.VENV_FAILED_CREATION"
 
@@ -62,18 +68,23 @@ def add_gitignore(_name):
     create_venv.add_gitignore = add_gitignore
 
     args = []
-    if git_ignore:
-        args.append("--git-ignore")
-    if install:
-        args.append("--install")
+    if git_ignore == "useGitIgnore":
+        args += ["--git-ignore"]
+    if install == "requirements":
+        args += ["--requirements", "requirements-for-test.txt"]
+    elif install == "toml":
+        args += ["--toml", "pyproject.toml", "--extras", "test"]
+
     create_venv.main(args)
-    assert install_packages_called == install
+    assert install_packages_called == (install != "skipInstall")
 
     # run_process is called when the venv does not exist
-    assert run_process_called != env_exists
+    assert run_process_called == (env_exists == "noEnv")
 
     # add_gitignore is called when new venv is created and git_ignore is True
-    assert add_gitignore_called == (not env_exists and git_ignore)
+    assert add_gitignore_called == (
+        (env_exists == "noEnv") and (git_ignore == "useGitIgnore")
+    )
 
 
 @pytest.mark.parametrize("install_type", ["requirements", "pyproject"])
@@ -93,12 +104,79 @@ def run_process(args, error_message):
         elif args[1:-1] == ["-m", "pip", "install", "-r"]:
             installing = "requirements"
             assert error_message == "CREATE_VENV.PIP_FAILED_INSTALL_REQUIREMENTS"
-        elif args[1:] == ["-m", "pip", "install", "-e", ".[extras]"]:
+        elif args[1:] == ["-m", "pip", "install", "-e", ".[test]"]:
             installing = "pyproject"
             assert error_message == "CREATE_VENV.PIP_FAILED_INSTALL_PYPROJECT"
 
     create_venv.run_process = run_process
 
-    create_venv.main(["--install"])
+    if install_type == "requirements":
+        create_venv.main(["--requirements", "requirements-for-test.txt"])
+    elif install_type == "pyproject":
+        create_venv.main(["--toml", "pyproject.toml", "--extras", "test"])
+
     assert pip_upgraded
     assert installing == install_type
+
+
+@pytest.mark.parametrize(
+    ("extras", "expected"),
+    [
+        ([], ["-m", "pip", "install", "-e", "."]),
+        (["test"], ["-m", "pip", "install", "-e", ".[test]"]),
+        (["test", "doc"], ["-m", "pip", "install", "-e", ".[test,doc]"]),
+    ],
+)
+def test_toml_args(extras, expected):
+    importlib.reload(create_venv)
+
+    actual = []
+
+    def run_process(args, error_message):
+        nonlocal actual
+        actual = args[1:]
+
+    create_venv.run_process = run_process
+
+    create_venv.install_toml(sys.executable, extras)
+
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    ("extras", "expected"),
+    [
+        ([], None),
+        (
+            ["requirements/test.txt"],
+            [sys.executable, "-m", "pip", "install", "-r", "requirements/test.txt"],
+        ),
+        (
+            ["requirements/test.txt", "requirements/doc.txt"],
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "-r",
+                "requirements/test.txt",
+                "-r",
+                "requirements/doc.txt",
+            ],
+        ),
+    ],
+)
+def test_requirements_args(extras, expected):
+    importlib.reload(create_venv)
+
+    actual = None
+
+    def run_process(args, error_message):
+        nonlocal actual
+        actual = args
+
+    create_venv.run_process = run_process
+
+    create_venv.install_requirements(sys.executable, extras)
+
+    assert actual == expected

pythonFiles/tests/test_create_venv.py

@@ -437,6 +437,8 @@ export namespace CreateEnv {
         export const selectPythonQuickPickTitle = l10n.t('Select a python to use for environment creation');
         export const providerDescription = l10n.t('Creates a `.venv` virtual environment in the current workspace');
         export const error = l10n.t('Creating virtual environment failed with error.');
+        export const tomlExtrasQuickPickTitle = l10n.t('Select optional dependencies to install from pyproject.toml');
+        export const requirementsQuickPickTitle = l10n.t('Select dependencies to install');
     }
 
     export namespace Conda {
@@ -454,13 +456,11 @@ export namespace CreateEnv {
 
 export namespace ToolsExtensions {
     export const flake8PromptMessage = l10n.t(
-        'toolsExt.flake8.message',
         'Use the Flake8 extension to enable easier configuration and new features such as quick fixes.',
     );
     export const pylintPromptMessage = l10n.t(
-        'toolsExt.pylint.message',
         'Use the Pylint extension to enable easier configuration and new features such as quick fixes.',
     );
-    export const installPylintExtension = l10n.t('toolsExt.install.pylint', 'Install Pylint extension');
-    export const installFlake8Extension = l10n.t('toolsExt.install.flake8', 'Install Flake8 extension');
+    export const installPylintExtension = l10n.t('Install Pylint extension');
+    export const installFlake8Extension = l10n.t('Install Flake8 extension');
 }

src/client/common/utils/localize.ts

@@ -1,7 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { ConfigurationScope, workspace, WorkspaceConfiguration, WorkspaceEdit, WorkspaceFolder } from 'vscode';
+import {
+    CancellationToken,
+    ConfigurationScope,
+    GlobPattern,
+    Uri,
+    workspace,
+    WorkspaceConfiguration,
+    WorkspaceEdit,
+    WorkspaceFolder,
+} from 'vscode';
 import { Resource } from '../types';
 
 export function getWorkspaceFolders(): readonly WorkspaceFolder[] | undefined {
@@ -23,3 +32,12 @@ export function getConfiguration(section?: string, scope?: ConfigurationScope |
 export function applyEdit(edit: WorkspaceEdit): Thenable<boolean> {
     return workspace.applyEdit(edit);
 }
+
+export function findFiles(
+    include: GlobPattern,
+    exclude?: GlobPattern | null,
+    maxResults?: number,
+    token?: CancellationToken,
+): Thenable<Uri[]> {
+    return workspace.findFiles(include, exclude, maxResults, token);
+}