python-mommy-venv/python_mommy_venv/__main__.py

import sys
from pathlib import Path
import stat
import subprocess
import logging
import json
import argparse

from .responses import compile_config
from .static import IS_VENV, VENV_DIRECTORY, CONFIG_DIRECTORY, COMPILED_CONFIG_FILE_NAME

logging.basicConfig(
    format='%(message)s',
    force=True,
)

log_level = logging.INFO

mommy_logger = logging.getLogger("mommy")
mommy_logger.setLevel(logging.INFO)
serious_logger = logging.getLogger("serious")
serious_logger.setLevel(50)


def config_logging(verbose: bool):
    if verbose:
        logging.basicConfig(
            format=logging.BASIC_FORMAT,
            force=True,
        )
        logging.getLogger().setLevel(logging.DEBUG)
        mommy_logger.setLevel(50)
        serious_logger.setLevel(logging.DEBUG)
    

WRAPPER_TEMPLATE = """#!{inner_bin}
# -*- coding: utf-8 -*-

import sys, subprocess
from python_mommy_venv import get_response


INTERPRETER = "{inner_bin}"
result = subprocess.run([INTERPRETER] + sys.argv[1:])
code = result.returncode

print()
print(get_response(code))
exit(code=code)
"""


PIP_HOOK = """# GENERATED BY MOMMY
    code = main()

    from pathlib import Path

    bin_path = Path(".venv", "bin")

    python_interpreter_wrappers = []
    for path in bin_path.iterdir():
        if path.is_symlink() and path.name.startswith("inner_"):
            python_interpreter_wrappers.append(path.name.replace("inner_", "", 1))

    for path in bin_path.iterdir():
        if path.is_symlink():
            continue
        if path.name in python_interpreter_wrappers:
            continue

        text: str
        with path.open("r") as f:
            text = f.read()

        first_line = text.split("\\n")[0]
        if not ("inner_" in first_line and first_line.startswith("#!")):
            continue 
        
        print(f"mommifying " + str(path))

        text = text.replace("inner_", "", 1)
        with path.open("w") as f:
            f.write(text)

    sys.exit(code)"""


def assert_venv(only_warn: bool = False):
    if not IS_VENV:
        mommy_logger.error("mommy doesn't run in a virtual environment~")
        serious_logger.error("this should run in a virtual environment")
        if not only_warn:
            exit(1)


def write_compile_config(local: bool):
    assert_venv(only_warn=not local)

    compiled_base_dir = VENV_DIRECTORY if local else CONFIG_DIRECTORY
    compiled_config_file = compiled_base_dir / COMPILED_CONFIG_FILE_NAME
    compiled = compile_config()
    mommy_logger.info("mommy writes its moods in %s", compiled_config_file)
    serious_logger.info("writing compiled config file to %s", compiled_config_file)
    compiled_base_dir.mkdir(parents=True, exist_ok=True)
    with compiled_config_file.open("w") as f:
        json.dump(compiled, f, indent=4)
    if not local:
        (VENV_DIRECTORY / COMPILED_CONFIG_FILE_NAME).unlink(missing_ok=True)


def wrap_interpreter(path: Path):
    mommy_logger.info("mommy found a symlink to an interpreter~ %s", str(path))
    serious_logger.info("interpreter symlink found at %s", str(path))

    inner_symlink = path.parent / ("inner_" + path.name)
    symlink_target = path.resolve()

    if inner_symlink.exists():
        raise Exception("inner symlink somehow already exists. This shouldn't happen because of prior checks")

    mommy_logger.info("mommy shows her girl where the interpreter is: %s -> %s", inner_symlink, symlink_target)
    serious_logger.info("creating symlink: %s -> %s", inner_symlink, symlink_target)
    inner_symlink.symlink_to(symlink_target)

    # remove original symlink
    mommy_logger.info("mommy deletes the original interpreter~ %s", path)
    serious_logger.info("deleting original symlink %s", path)
    path.unlink()

    # creating the wrapper string
    mommy_logger.info("mommy writes wrapper script as %s", Path)
    serious_logger.info("writing wrapper script at %s", path)
    with path.open("w") as f:
        f.write(WRAPPER_TEMPLATE.format(inner_bin=str(inner_symlink)))
    serious_logger.info("making wrapper script executable")
    path.chmod(path.stat().st_mode | stat.S_IEXEC)


def install_pip_hook(path: Path):
    text: str
    with path.open("r") as f:
        text = f.read()

    if "# GENERATED BY MOMMY" in text:
        mommy_logger.info("ahhhhh mommy already watches %s", str(path))
        serious_logger.info("pip hook already installed at %s", str(path))
        return
    
    mommy_logger.info("mommy needs to keep an eye on this little pip~ %s", str(path))
    serious_logger.info("installing pip hook at %s", str(path))

    text = text.replace("sys.exit(main())", PIP_HOOK, 1)
    with path.open("w") as f:
        f.write(text)


def cli_compile_config():
    parser = argparse.ArgumentParser(description="only recompile the config")
    
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="enable verbose and serious output"
    )

    parser.add_argument(
        "-l", "--local",
        action="store_true",
        help="compile the config only for the current virtual environment"
    )

    args = parser.parse_args()

    write_compile_config(args.local)


def mommify_venv():
    parser = argparse.ArgumentParser(description="patch the virtual environment to use mommy")
    
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="enable verbose and serious output"
    )

    parser.add_argument(
        "-l", "--local",
        action="store_true",
        help="compile the config only for the current virtual environment"
    )

    args = parser.parse_args()

    config_logging(args.verbose)
    assert_venv()

    write_compile_config(args.local)

    mommy_logger.info("")

    bin_path = VENV_DIRECTORY / "bin"
    bin_path = bin_path.resolve()

    mommy_logger.info("mommy looks in %s to mess your system up~ <33", str(bin_path))
    serious_logger.info("scanning binary directory of venv at %s", str(bin_path))

    for path in list(bin_path.iterdir()):
        name = path.name

        if path.is_symlink():
            # could be python interpreter
            # check for both just to be more expressive
            if name.startswith("inner_"):
                continue
            
            if subprocess.run([str(path), '-c', '"exit(0)"']).returncode != 0:
                continue

            wrap_interpreter(path)

        else:
            # could be pip
            if not name.startswith("pip"):
                continue

            install_pip_hook(path)