Added several virtualenv improvements.

Fixes #25
Fixes #26
Fixes #27

Signed-off-by: Pedro Algarvio <[email protected]>

Author: s0undt3ch

changelog/25.improvement.rst

@@ -0,0 +1,2 @@
+Allow importing the virtualenv's dependencies into the python running ``tools``.
+This will allow maintaining a lighter ``tools.txt`` requirements file, and install additional dependencies only for the commands that really need them.

changelog/26.improvement.rst

@@ -0,0 +1 @@
+Add support for a default Virtualenv. The default virtualenv site-packages will be added to the running python as an extra site dir.

changelog/27.improvement.rst

@@ -0,0 +1 @@
+Allow passing virtualenv configuration when calling ``ptscripts.register_tools_module``

src/ptscripts/__init__.py

@@ -6,13 +6,24 @@
 
 import ptscripts.logs
 from ptscripts.parser import command_group
-from ptscripts.parser import Context, RegisteredImports
+from ptscripts.parser import Context, RegisteredImports, DefaultVirtualenvConfig
+from ptscripts.virtualenv import VirtualEnvConfig
 
 __all__ = ["command_group", "register_tools_module", "Context", "CWD"]
 
 
-def register_tools_module(import_module: str) -> None:
+def register_tools_module(import_module: str, venv_config: VirtualEnvConfig | None = None) -> None:
     """
     Register a module to be imported when instantiating the tools parser.
     """
-    RegisteredImports.register_import(import_module)
+    RegisteredImports.register_import(import_module, venv_config=venv_config)
+
+
+def set_default_venv_config(venv_config: VirtualEnvConfig) -> None:
+    """
+    Define the default virtualenv configuration.
+
+    This virtualenv will be available to all commands, and it's ``site-packages``
+    dir(s) will be added to the current python interpreter site.
+    """
+    DefaultVirtualenvConfig.set_default_venv_config(venv_config)

src/ptscripts/parser.py

@@ -14,12 +14,14 @@
 import typing
 from collections.abc import Iterator
 from contextlib import contextmanager
+from contextlib import nullcontext
 from functools import partial
 from subprocess import CompletedProcess
 from types import FunctionType
 from types import GenericAlias
 from typing import Any
 from typing import cast
+from typing import ContextManager
 from typing import TYPE_CHECKING
 from typing import TypedDict
 
@@ -248,38 +250,70 @@ def web(self) -> requests.Session:
         return requests.Session()
 
 
+class DefaultVirtualenvConfig:
+    """
+    Simple class to hold registered imports.
+    """
+
+    _instance: DefaultVirtualenvConfig | None = None
+    venv_config: VirtualEnvConfig
+
+    def __new__(cls):
+        """
+        Method that instantiates a singleton class and returns it.
+        """
+        if cls._instance is None:
+            instance = super().__new__(cls)
+            cls._instance = instance
+        return cls._instance
+
+    @classmethod
+    def set_default_venv_config(cls, venv_config: VirtualEnvConfig) -> None:
+        """
+        Register an import.
+        """
+        instance = cls._instance
+        if instance is None:
+            instance = cls()
+        if venv_config and "name" not in venv_config:
+            venv_config["name"] = "default"
+        instance.venv_config = venv_config
+
+
 class RegisteredImports:
     """
     Simple class to hold registered imports.
     """
 
     _instance: RegisteredImports | None = None
-    _registered_imports: list[str]
+    _registered_imports: dict[str, VirtualEnvConfig | None]
 
     def __new__(cls):
         """
         Method that instantiates a singleton class and returns it.
         """
         if cls._instance is None:
             instance = super().__new__(cls)
-            instance._registered_imports = []
+            instance._registered_imports = {}
             cls._instance = instance
         return cls._instance
 
     @classmethod
-    def register_import(cls, import_module: str) -> None:
+    def register_import(
+        cls, import_module: str, venv_config: VirtualEnvConfig | None = None
+    ) -> None:
         """
         Register an import.
         """
         instance = cls()
         if import_module not in instance._registered_imports:
-            instance._registered_imports.append(import_module)
+            instance._registered_imports[import_module] = venv_config
 
     def __iter__(self):
         """
         Return an iterator of all registered imports.
         """
-        return iter(self._registered_imports)
+        return iter(self._registered_imports.items())
 
 
 class Parser:
@@ -370,14 +404,29 @@ def __new__(cls):
         return cls._instance
 
     def _process_registered_tool_modules(self):
-        for module_name in RegisteredImports():
-            try:
-                importlib.import_module(module_name)
-            except ImportError as exc:
-                if os.environ.get("TOOLS_IGNORE_IMPORT_ERRORS", "0") == "0":
-                    self.context.warn(
-                        f"Could not import the registered tools module {module_name!r}: {exc}"
-                    )
+        default_venv: VirtualEnv | ContextManager[None]
+        default_venv_config = DefaultVirtualenvConfig().venv_config
+        if default_venv_config:
+            default_venv = VirtualEnv(ctx=self.context, **default_venv_config)
+        else:
+            default_venv = nullcontext()
+        with default_venv:
+            for module_name, venv_config in RegisteredImports():
+                venv: VirtualEnv | ContextManager[None]
+                if venv_config:
+                    if "name" not in venv_config:
+                        venv_config["name"] = module_name
+                    venv = VirtualEnv(ctx=self.context, **venv_config)
+                else:
+                    venv = nullcontext()
+                with venv:
+                    try:
+                        importlib.import_module(module_name)
+                    except ImportError as exc:
+                        if os.environ.get("TOOLS_IGNORE_IMPORT_ERRORS", "0") == "0":
+                            self.context.warn(
+                                f"Could not import the registered tools module {module_name!r}: {exc}"
+                            )
 
     def parse_args(self):
         """
@@ -471,6 +520,8 @@ def __init__(self, name, help, description=None, parent=None, venv_config=None):
             GroupReference.add_command(tuple(parent + [name]), self)
             parent = GroupReference()[tuple(parent)]
 
+        if venv_config and "name" not in venv_config:
+            venv_config["name"] = self.name
         self.venv_config = venv_config or {}
         self.parser = parent.subparsers.add_parser(
             name.replace("_", "-"),
@@ -634,22 +685,22 @@ def __call__(self, func, options, venv_config: VirtualEnvConfig | None = None):
                     kwargs[name] = getattr(options, name)
 
         bound = signature.bind_partial(*args, **kwargs)
-        venv = None
+        venv: VirtualEnv | ContextManager[None]
         if venv_config:
-            venv_name = getattr(options, f"{self.name}_command")
-            venv = VirtualEnv(name=f"{self.name}.{venv_name}", ctx=self.context, **venv_config)
+            if "name" not in venv_config:
+                venv_config["name"] = getattr(options, f"{self.name}_command")
+            venv = VirtualEnv(ctx=self.context, **venv_config)
         elif self.venv_config:
-            venv = VirtualEnv(name=self.name, ctx=self.context, **self.venv_config)
-        if venv:
-            with venv:
-                previous_venv = self.context.venv
-                try:
-                    self.context.venv = venv
-                    func(self.context, *bound.args, **bound.kwargs)
-                finally:
-                    self.context.venv = previous_venv
+            venv = VirtualEnv(ctx=self.context, **self.venv_config)
         else:
-            func(self.context, *bound.args, **bound.kwargs)
+            venv = nullcontext()
+        with venv:
+            previous_venv = self.context.venv
+            try:
+                self.context.venv = venv
+                func(self.context, *bound.args, **bound.kwargs)
+            finally:
+                self.context.venv = previous_venv
 
 
 def command_group(

src/ptscripts/virtualenv.py

@@ -6,6 +6,7 @@
 import os
 import pathlib
 import shutil
+import site
 import subprocess
 import sys
 import textwrap
@@ -32,12 +33,14 @@ class VirtualEnvConfig(TypedDict):
     Virtualenv Configuration Typing.
     """
 
+    name: NotRequired[str]
     requirements: NotRequired[list[str]]
     requirements_files: NotRequired[list[pathlib.Path]]
     env: NotRequired[dict[str, str]]
     system_site_packages: NotRequired[bool]
     pip_requirement: NotRequired[str]
     setuptools_requirement: NotRequired[str]
+    add_as_extra_site_packages: NotRequired[bool]
 
 
 def _cast_to_pathlib_path(value):
@@ -60,6 +63,7 @@ class VirtualEnv:
     system_site_packages: bool = attr.ib(default=False)
     pip_requirement: str = attr.ib(repr=False)
     setuptools_requirement: str = attr.ib(repr=False)
+    add_as_extra_site_packages: bool = attr.ib(default=False)
     environ: dict[str, str] = attr.ib(init=False, repr=False)
     venv_dir: pathlib.Path = attr.ib(init=False)
     venv_python: pathlib.Path = attr.ib(init=False, repr=False)
@@ -175,6 +179,25 @@ def _create_virtualenv(self):
             self.setuptools_requirement,
         )
 
+    def _add_as_extra_site_packages(self):
+        if self.add_as_extra_site_packages is False:
+            return
+        ret = self.run_code(
+            "import json,site; print(json.dumps(site.getsitepackages()))",
+            capture=True,
+            check=False,
+        )
+        if ret.returncode:
+            self.ctx.error(
+                f"Failed to get the virtualenv's site packages path: {ret.stderr.decode()}"
+            )
+            self.ctx.exit(1)
+        site_packages = site.getsitepackages()
+        for path in json.loads(ret.stdout.strip().decode()):
+            if path not in site_packages:
+                site.addsitedir(path)
+        site_packages = site.getsitepackages()
+
     def __enter__(self):
         """
         Creates the virtual environment when entering context.
@@ -184,6 +207,7 @@ def __enter__(self):
         except subprocess.CalledProcessError:
             raise AssertionError("Failed to create virtualenv")
         self._install_requirements()
+        self._add_as_extra_site_packages()
         return self
 
     def __exit__(self, *args):