Mojo venv
import click
from click.testing import CliRunner


@click.command()
@click.option(
    "--without-scm-ignore-files",
    "scm_ignore_files",
    is_flag=True,
    type=frozenset,
    flag_value=frozenset(),
    default=frozenset(["git"]),
    help="Skips adding SCM ignore files to the environment "
    "directory (Git is supported by default).",
)
def rcli(scm_ignore_files):
    print(f"{scm_ignore_files = }")
    print(f"{type(scm_ignore_files) = }")


def test_rcli():
    runner = CliRunner()
    result = runner.invoke(rcli)
    assert result.exit_code == 0
    print(result.output)


if __name__ == "__main__":
    test_rcli()

notes/reproducer.py
mojopaa/menv
import configparser
import logging
import os
import shutil
import subprocess
import sys
import types
from pathlib import Path

import tomlkit
from tomlkit.toml_file import TOMLFile

MODULAR_NAME = ".modular"
MODULAR_PKG_FOLDER = "pkg"
MODULAR_PKG_NAME = "packages.modular.com_mojo"
MODULAR_CONFIG_NAME = "modular.cfg"


MODULAR_DIR = Path.home() / MODULAR_NAME
MODULAR_CONFIG = MODULAR_DIR / MODULAR_CONFIG_NAME
MOJO_PKG_DIR = MODULAR_DIR / MODULAR_PKG_FOLDER / MODULAR_PKG_NAME
MOJO_BIN_DIR = MOJO_PKG_DIR / "bin"
MOJO_LIB_DIR = MOJO_PKG_DIR / "lib"
MOJO_IMPORT_DIR = MOJO_LIB_DIR / "mojo"
MOJO_EXECUTABLE = MOJO_BIN_DIR / "mojo"

logger = logging.getLogger(__name__)


def create_if_needed(d):
    if not os.path.exists(d):
        os.makedirs(d)
    elif os.path.islink(d) or os.path.isfile(d):
        raise ValueError("Unable to create directory %r" % d)


def clear_directory(path):
    for fn in os.listdir(path):
        fn = os.path.join(path, fn)
        if os.path.islink(fn) or os.path.isfile(fn):
            os.remove(fn)
        elif os.path.isdir(fn):
            shutil.rmtree(fn)


def change_config(path: str, section: str, key: str, value: str):
    config = configparser.ConfigParser()
    config.read(path)

    config[section][key] = value

    with open(path, "w", encoding="utf-8") as f:
        config.write(f)


class MojoEnvBuilder:
    def __init__(
        self,
        system_site_packages=False,
        clear=False,
        symlinks=False,
        upgrade=False,
        prompt=None,
        upgrade_deps=False,
        scm_ignore_files=True,
    ):
        self.system_site_packages = system_site_packages
        self.clear = clear
        self.symlinks = symlinks
        self.upgrade = upgrade
        self.orig_prompt = prompt
        if prompt == ".":  # see bpo-38901
            prompt = os.path.basename(os.getcwd())
        self.prompt = prompt
        self.upgrade_deps = upgrade_deps
        self.scm_ignore_files = scm_ignore_files

    def create(self, env_dir):
        """
        Create a virtual environment in a directory.

        :param env_dir: The target directory to create an environment in.

        """
        env_dir = os.path.abspath(env_dir)
        context = self.ensure_directories(env_dir)
        # for scm in self.scm_ignore_files:
        #     getattr(self, f"create_{scm}_ignore_file")(context)
        # See issue 24875. We need system_site_packages to be False
        # until after pip is installed.
        true_system_site_packages = self.system_site_packages
        self.system_site_packages = False
        self.create_configuration(context)
        self.setup_mojo(context)
        # if self.with_pip:
        #     self._setup_pip(context)
        if not self.upgrade:
            self.setup_scripts(context)
            self.post_setup(context)
        if true_system_site_packages:
            # We had set it to False before, now
            # restore it and rewrite the configuration
            self.system_site_packages = True
            self.create_configuration(context)
        if self.upgrade_deps:
            self.upgrade_dependencies(context)
        if self.scm_ignore_files:
            self.create_git_ignore_file(context)

    def ensure_directories(self, env_dir: str | Path):
        """
        Create the directories for the environment.

        Returns a context object which holds paths in the environment,
        for use by subsequent logic.
        """
        if os.pathsep in os.fspath(env_dir):
            raise ValueError(
                f"Refusing to create a venv in {env_dir} because "
                f"it contains the PATH separator {os.pathsep}."
            )

        if os.path.exists(env_dir) and self.clear:
            clear_directory(env_dir)

        context = types.SimpleNamespace()
        context.env_dir = str(env_dir)
        context.env_name = os.path.split(env_dir)[1]
        prompt = self.prompt if self.prompt is not None else context.env_name
        context.prompt = "(%s) " % prompt
        create_if_needed(env_dir)

        context.mojo_dir = str(MOJO_BIN_DIR)  # TODO: use findmojo
        context.mojo_exe = "mojo"

        # venv paths
        venv_modular_dir = Path(env_dir) / MODULAR_NAME
        venv_modular_cfg = venv_modular_dir / MODULAR_CONFIG_NAME
        venv_pkg_dir = venv_modular_dir / MODULAR_PKG_FOLDER / MODULAR_PKG_NAME

        bin_name = "bin"  # Mojo bin name
        venv_bin_dir = venv_pkg_dir / bin_name
        venv_lib_dir = venv_pkg_dir / "lib"
        venv_mojo_excutable = venv_bin_dir / "mojo"

        context.executable = MOJO_EXECUTABLE
        context.bin_name = bin_name
        context.pkg_dir = str(venv_pkg_dir)
        context.bin_path = str(venv_bin_dir)
        context.lib_path = str(venv_lib_dir)
        context.env_exe = str(venv_mojo_excutable)
        context.env_cfg = str(venv_modular_cfg)

        # venv bin path, reference: /usr/lib/python3.10/venv/__init__.py
        if sys.platform == "win32":
            py_venv_biname = "Scripts"
        else:
            py_venv_biname = "bin"

        context.py_venv_binpath = str(Path(env_dir) / py_venv_biname)

        create_if_needed(venv_bin_dir)
        create_if_needed(venv_lib_dir)

        return context

    def create_configuration(self, context):
        """
        Create a configuration file indicating where the environment's Python
        was copied from, and whether the system site-packages should be made
        available in the environment.

        Args:
            context: The context object containing information about the environment.

        Returns:
            None
        """
        context.cfg_path = path = os.path.join(context.env_dir, "mojovenv.toml")

        cfg = tomlkit.document()
        mojo_tab = tomlkit.table()
        mojo_tab.add("home", context.mojo_dir)

        if self.system_site_packages:
            incl = True
        else:
            incl = False
        mojo_tab.add("include-system-site-packages", incl)

        # Read VERSION
        with open(MOJO_PKG_DIR / "VERSION") as f:
            version = f.read().strip()

        mojo_tab.add("version", version)

        if self.prompt is not None:
            mojo_tab.add("prompt", f"{self.prompt!r}")

        mojo_tab.add("mojo-executable", str(MOJO_EXECUTABLE))

        # build command
        args = []
        nt = os.name == "nt"
        if nt and self.symlinks:
            args.append("--symlinks")
        if not nt and not self.symlinks:
            args.append("--copies")
        if self.system_site_packages:
            args.append("--system-site-packages")
        if self.clear:
            args.append("--clear")
        if self.upgrade:
            args.append("--upgrade")
        if self.upgrade_deps:
            args.append("--upgrade-deps")
        if self.orig_prompt is not None:
            args.append(f'--prompt="{self.orig_prompt}"')
        if not self.scm_ignore_files:
            args.append("--without-scm-ignore-files")

        args.append(context.env_dir)
        args = " ".join(args)

        mojo_tab.add("command", f"menv {args}")
        cfg["mojo"] = mojo_tab
        TOMLFile(path).write(cfg)

    if os.name != "nt":

        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
            """
            Try symlinking a file, and if that fails, fall back to copying.
            """
            force_copy = not self.symlinks
            if not force_copy:
                try:
                    if not os.path.islink(dst):  # can't link to itself!
                        if relative_symlinks_ok:
                            assert os.path.dirname(src) == os.path.dirname(dst)
                            os.symlink(os.path.basename(src), dst)
                        else:
                            os.symlink(src, dst)
                except Exception:  # may need to use a more specific exception
                    logger.warning("Unable to symlink %r to %r", src, dst)
                    force_copy = True
            if force_copy:
                shutil.copyfile(src, dst)

    else:

        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
            """
            Try symlinking a file, and if that fails, fall back to copying.
            """
            bad_src = os.path.lexists(src) and not os.path.exists(src)
            if self.symlinks and not bad_src and not os.path.islink(dst):
                try:
                    if relative_symlinks_ok:
                        assert os.path.dirname(src) == os.path.dirname(dst)
                        os.symlink(os.path.basename(src), dst)
                    else:
                        os.symlink(src, dst)
                    return
                except Exception:  # may need to use a more specific exception
                    logger.warning("Unable to symlink %r to %r", src, dst)

            if not os.path.exists(src):
                if not bad_src:
                    logger.warning("Unable to copy %r", src)
                return

            shutil.copyfile(src, dst)

    def recursive_symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
        for item in os.listdir(src):
            src_item = os.path.join(src, item)
            dst_item = os.path.join(dst, item)
            if os.path.isfile(src_item):
                self.symlink_or_copy(src_item, dst_item, relative_symlinks_ok)
                shutil.copymode(src_item, dst_item)
            elif os.path.isdir(src_item):
                os.makedirs(dst_item, exist_ok=True)
                self.recursive_symlink_or_copy(src_item, dst_item, relative_symlinks_ok)
            else:
                logger.warning(f"Skipping {src_item}")

    def create_git_ignore_file(self, context):
        """
        Create a .gitignore file in the environment directory.

        The contents of the file cause the entire environment directory to be
        ignored by git.
        """
        gitignore_path = os.path.join(context.env_dir, ".gitignore")
        with open(gitignore_path, "w", encoding="utf-8") as file:
            file.write(
                "# Created by venv; "
                "see https://docs.python.org/3/library/venv.html\n"
            )  # TODO: change this?
            file.write("*\n")

    def setup_mojo(self, context):
        """
        Set up a Mojo executable in the environment.

        Args:
            context (obj): The information for the environment creation request being processed.
        """
        binpath = context.bin_path
        libpath = context.lib_path
        path = context.env_exe
        copier = self.symlink_or_copy
        dirname = context.mojo_dir  # context.mojo_dir = str(MOJO_BIN_DIR)

        if os.name != "nt":
            # copy lib and bin to venv
            self.recursive_symlink_or_copy(MOJO_LIB_DIR, libpath)
            self.recursive_symlink_or_copy(MOJO_BIN_DIR, binpath)

            for bin_item in os.listdir(binpath):
                if not os.path.islink(os.path.join(binpath, bin_item)):
                    # Set the executable's permissions
                    os.chmod(os.path.join(binpath, bin_item), 0o755)

            # Create symbolic links for mojo executables
            for suffix in "mojo":
                path = os.path.join(binpath, suffix)
                if not os.path.exists(path):
                    # Make copies if symlinks are not wanted
                    copier(context.env_exe, path, relative_symlinks_ok=True)
                    if not os.path.islink(path):
                        os.chmod(path, 0o755)

            # Copy config and change import_path settings etc.
            copier(MODULAR_CONFIG, context.env_cfg)
            self.write_modular_cfg(context)

        else:
            pass  # TODO

    def write_modular_cfg(self, context):
        cfg_path = context.env_cfg  # context.env_cfg = str(venv_modular_cfg)
        libpath = context.lib_path  # str(venv_pkg_dir / "lib")

        change_config(
            cfg_path,
            section="mojo",
            key="import_path",
            value=os.path.join(libpath, "mojo"),
        )  # not working
        change_config(
            cfg_path,
            section="installed",
            key="packages_modular_com_mojo",
            value=context.pkg_dir,
        )
        # venv_pkg_config = os.path.join(context.pkg_dir, MODULAR_CONFIG_NAME)
        # print(venv_pkg_config)
        # change_config(venv_pkg_config, section="mojo", key="import_path", value=os.path.join(libpath, "mojo"))

    def replace_variables(self, text, context):
        """
        Replace variable placeholders in the script text with context-specific
        variables.

        Args:
            text (str): The text in which to replace placeholder variables.
            context (Context): The information for the environment creation request
                being processed.

        Returns:
            str: The text passed in, but with variables replaced.
        """
        # Replace '__VENV_DIR__' placeholder with context.env_dir
        text = text.replace("__VENV_DIR__", context.env_dir)  # abspath of venv

        # Replace '__VENV_NAME__' placeholder with context.env_name
        text = text.replace("__VENV_NAME__", context.env_name)  # stem of venv

        # Replace '__VENV_PROMPT__' placeholder with context.prompt
        text = text.replace("__VENV_PROMPT__", context.prompt)

        # Replace '__VENV_BIN_NAME__' placeholder with context.bin_name
        text = text.replace("__VENV_BIN_NAME__", context.bin_name)  # bin

        # Replace '__VENV_BIN_PATH__' placeholder with context.bin_path
        text = text.replace("__VENV_BIN_PATH__", context.bin_path)

        # Replace '__VENV_MOJO__' placeholder with context.env_exe
        text = text.replace("__VENV_MOJO__", context.env_exe)

        return text

    def setup_scripts(self, context):
        """
        Set up scripts into the created environment from a directory.

        This method installs the default scripts into the environment
        being created. You can prevent the default installation by overriding
        this method if you really need to, or if you need to specify
        a different location for the scripts to install. By default, the
        'scripts' directory in the venv package is used as the source of
        scripts to install.
        """
        path = str(Path(__file__).with_name("scripts").absolute())
        self.install_scripts(context, path)

    def install_scripts(self, context, path):
        """
        Install scripts into the created environment from a directory.

        Args:
            context: The information for the environment creation request
                being processed.
            path: Absolute pathname of a directory containing script.
                Scripts in the 'common' subdirectory of this directory,
                and those in the directory named for the platform
                being run on, are installed in the created environment.
                Placeholder variables are replaced with environment-
                specific values.
        """
        # binpath = context.bin_path  # Get the bin path from the context
        binpath = context.py_venv_binpath  # Get the bin path from the context
        plen = len(path)  # Get the length of the path
        for root, dirs, files in os.walk(path):
            if root == path:  # At top-level, remove irrelevant dirs
                for d in dirs[:]:
                    if d not in ("common", os.name):
                        dirs.remove(d)
                continue  # Ignore files in top level
            for f in files:
                if (
                    os.name == "nt"
                    and f.startswith("python")
                    and f.endswith((".exe", ".pdb"))
                ):
                    continue  # Skip files that start with 'python' and end with '.exe' or '.pdb' on Windows
                srcfile = os.path.join(root, f)  # Get the source file path
                suffix = root[plen:].split(os.sep)[
                    2:
                ]  # Get the relative path of the file
                if not suffix:
                    dstdir = binpath  # If there is no suffix, set the destination directory as binpath
                else:
                    dstdir = os.path.join(
                        binpath, *suffix
                    )  # Set the destination directory as binpath + suffix
                if not os.path.exists(dstdir):
                    os.makedirs(
                        dstdir
                    )  # Create the destination directory if it does not exist
                dstfile = os.path.join(dstdir, f)  # Get the destination file path
                with open(srcfile, "rb") as f:
                    data = f.read()  # Read the source file data
                if not srcfile.endswith((".exe", ".pdb")):
                    try:
                        data = data.decode("utf-8")  # Decode the data as utf-8
                        data = self.replace_variables(
                            data, context
                        )  # Replace the placeholder variables with environment-specific values
                        data = data.encode("utf-8")  # Encode the data as utf-8
                    except UnicodeError as e:
                        data = None
                        logger.error(f"UnicodeError: {e}")
                        logger.warning(
                            "unable to copy script %r, " "may be binary: %s", srcfile, e
                        )  # Log a warning if unable to copy script due to UnicodeError
                if data is not None:
                    with open(dstfile, "wb") as f:
                        f.write(data)  # Write the data to the destination file
                    shutil.copymode(
                        srcfile, dstfile
                    )  # Copy the permissions from the source file to the destination file

    def upgrade_dependencies(self, context):
        logger.warning("TODO: upgrade CORE_VENV_DEPS")
        logger.warning("NOP right now.")

    def post_setup(self, context):
        """
        Hook for post-setup modification of the venv.
        Subclasses may install additional packages or scripts here,
        add activation shell scripts, etc.

        Parameters:
          - context: The information for the environment creation request
                     being processed.

        No changes are made in the function body.
        """

        pass


def create(
    env_dir,
    system_site_packages=False,
    clear=False,
    symlinks=False,
    prompt=None,
    upgrade_deps=False,
    *,
    scm_ignore_files=True,
):
    """Create a virtual environment in a directory."""
    builder = MojoEnvBuilder(
        system_site_packages=system_site_packages,
        clear=clear,
        symlinks=symlinks,
        prompt=prompt,
        upgrade_deps=upgrade_deps,
        scm_ignore_files=scm_ignore_files,
    )
    builder.create(env_dir)


if __name__ == "__main__":
    # m = MojoEnvBuilder()
    # c = m.ensure_directories(".asdf")
    # m.create_configuration(c)
    # m.setup_mojo(c)
    # m.create_git_ignore_file(c)
    create(".asdf")

src/menv/builder.py
mojopaa/menv
import os
from venv import EnvBuilder

import click

from .builder import MojoEnvBuilder
from .utils import CORE_VENV_DEPS

if os.name == "nt":
    use_symlinks = False
else:
    use_symlinks = True


@click.command(context_settings=dict(help_option_names=["-h", "--help"]))
@click.argument("dirs", nargs=-1)
@click.option(
    "--system-site-packages",
    "system_site",
    is_flag=True,
    help="Give the virtual environment access to the " "system site-packages dir.",
)
@click.option(
    "--symlinks",
    is_flag=True,
    default=use_symlinks,
    help="Try to use symlinks rather than copies, "
    "when symlinks are not the default for "
    "the platform.",
)
@click.option(
    "--copies",
    is_flag=True,
    help="Try to use copies rather than symlinks, "
    "even when symlinks are the default for "
    "the platform.",
)
@click.option(
    "--clear",
    is_flag=True,
    help="Delete the contents of the "
    "environment directory if it "
    "already exists, before "
    "environment creation.",
)
@click.option(
    "--upgrade",
    is_flag=True,
    help="Upgrade the environment "
    "directory to use this version "
    "of Python, assuming Python "
    "has been upgraded in-place.",
)
@click.option(
    "--without-pip",
    "with_pip",
    is_flag=True,
    default=True,
    flag_value=False,
    help="Skips installing or upgrading pip in the "
    "virtual environment (pip is bootstrapped "
    "by default)",
)
@click.option(
    "--prompt", help="Provides an alternative prompt prefix for " "this environment."
)
@click.option(
    "--upgrade-deps",
    is_flag=True,
    help=f'Upgrade core dependencies ({", ".join(CORE_VENV_DEPS)}) '
    "to the latest version in PyPI",
)
@click.option(
    "--without-scm-ignore-files",
    "scm_ignore_files",
    is_flag=True,
    flag_value=False,
    default=True,
    help="Skips adding SCM ignore files to the environment "
    "directory (Git is supported by default).",
)
def cli(
    dirs,
    system_site,
    symlinks,
    copies,
    clear,
    upgrade,
    with_pip,
    prompt,
    upgrade_deps,
    scm_ignore_files,
):
    if upgrade and clear:
        raise ValueError("you cannot supply --upgrade and --clear together.")

    if copies:
        # print(f"{copies = }")
        symlinks = False
    # print(f"{dir = }, {system_site = }, {symlinks = }, {clear = }, {upgrade = }, {with_pip = }, {prompt = }, {upgrade_deps = }")
    # defaults: dir = '.asdf', system_site = False, symlinks = False,
    # clear = False, upgrade = False, with_pip = True, prompt = None, upgrade_deps = False
    for d in dirs:
        py_venv_builder = EnvBuilder(
            system_site_packages=system_site,
            clear=clear,
            symlinks=symlinks,
            upgrade=upgrade,
            with_pip=with_pip,
            prompt=prompt,
            upgrade_deps=upgrade_deps,
        )

        py_venv_builder.create(d)

        # print(f"{scm_ignore_files = }")
        # if isinstance(scm_ignore_files, str):
        #     scm_ignore_files = eval(scm_ignore_files)
        mojo_venv_builder = MojoEnvBuilder(
            system_site_packages=system_site,
            clear=clear,
            symlinks=symlinks,
            upgrade=upgrade,
            prompt=prompt,
            upgrade_deps=upgrade_deps,
            scm_ignore_files=scm_ignore_files,
        )
        mojo_venv_builder.create(d)

src/menv/cli.py
mojopaa/menv
CORE_VENV_DEPS = ("pip",)  # TODO: assure

src/menv/utils.py
mojopaa/menv
import shutil
from pathlib import Path
from pprint import pprint

import pytest

from menv.builder import (
    MojoEnvBuilder,
    change_config,
    clear_directory,
    create_if_needed,
)


class TestMojoEnvBuilder:
    @classmethod
    def setup_class(cls):
        cls.builder = MojoEnvBuilder()
        cls.venv_path = Path(__file__).parent / "venv"
        cls.context = cls.builder.ensure_directories(cls.venv_path)

    def test_init(self) -> None:
        pprint(f"{self.context = }")  # use pytest -s to show this print
        pass

    def test_ensure_directories(self) -> None:
        assert hasattr(self.context, "env_dir")
        assert hasattr(self.context, "env_name")
        assert hasattr(self.context, "prompt")
        assert hasattr(self.context, "mojo_dir")
        assert hasattr(self.context, "mojo_exe")
        assert hasattr(self.context, "executable")
        assert hasattr(self.context, "bin_name")
        assert hasattr(self.context, "pkg_dir")
        assert hasattr(self.context, "bin_path")
        assert hasattr(self.context, "lib_path")
        assert hasattr(self.context, "env_exe")
        assert hasattr(self.context, "env_cfg")

    def test_install_scripts(self):
        self.builder.setup_scripts(self.context)

        assert Path(self.context.bin_path).joinpath("mactivate").exists()

    # comment to debug venv
    @classmethod
    def teardown_class(cls) -> None:
        shutil.rmtree(cls.venv_path)

tests/unit/test_builder.py
mojopaa/menv