Small python script to setup an virtualenv and keep it up to date.

setup_venv.py

"""
This script creates a virtual environment and updates its dependencies.

The following features are supported.
- The script checks if the venv exists. If it does it will not recreate it.
- The script checks if the requirements.txt has changed. If it has, packages of the venv will be updated.
- The script checks if itself has changed. If it has, the venv will be recreated.
"""

import os
import subprocess
import sys
import argparse
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

REQUIREMENTS_TIMESTAMP_FILENAME = ".requirements_timestamp"
SCRIPT_TIMESTAMP_FILENAME = ".setup_venv_timestamp"
PYTHON_CMD = "python"
SCRIPT_PATH = os.path.realpath(__file__)
SCRIPT_DIR = os.path.dirname(SCRIPT_PATH)
REQUIREMENTS_FILE = os.path.join(SCRIPT_DIR, "requirements.txt")

def get_file_timestamp(filepath):
    return os.path.getmtime(filepath)


def create_virtualenv(venv_dir, clear):
    """Creates a virtual environment with venv

    Parameters
    ----------
    venv_dir : path to the directory where the venv will be created
    clear    : if true the former venv will be deleted before creation of a new one.
    """
    sargs = [PYTHON_CMD]

    sargs += ["-m", "venv"]
    if clear:
        sargs += ["--clear"]

    sargs += [venv_dir]
    subprocess.run(sargs)
    python_executable = os.path.join(venv_dir, "Scripts", "python")

    # Upgrade pip to its newest version - does not work in some build environments
    subprocess.run([python_executable, "-m", "pip", "install", "--upgrade", "pip"])


def install_requirements(venv_dir, requirements_file):
    # Get the current working directory
    cwd = os.getcwd()
    logger.debug(f'install_requirements: {{old_cwd: "{cwd}"}}')
    python_executable = os.path.join(venv_dir, "Scripts", "python")

    # !!This is a hack!! The working directory is changed to the directory where this
    # script is contained to support relative imports in the requirements.txt.
    # The script and requiremnts.txt must therefore be always in the same directory.
    os.chdir(SCRIPT_DIR)
    cwd = os.getcwd()
    logger.debug(f'changed working directory to: {{cwd: "{cwd}"}}')
    logger.debug(f'install_requirements: {{python_loc: "{python_executable}"}}')
    subprocess.run([python_executable, "-m", "pip", "install", "-r", requirements_file])
    os.chdir(cwd)
    cwd = os.getcwd()
    logger.debug(f'changed working directory to: {{cwd: "{cwd}"}}')


def update_virtual_env(
    venv_dir,
    requirements_timestamp_file,
    current_requirements_timestamp,
    script_timestamp_file,
    current_script_timestamp,
):
    print("Installing/Updating requirements...")
    install_requirements(venv_dir, REQUIREMENTS_FILE)
    Path(requirements_timestamp_file).write_text(str(current_requirements_timestamp))
    Path(script_timestamp_file).write_text(str(current_script_timestamp))


def is_requirement_txt_outdated(
    current_requirements_timestamp, stored_requirements_timestamp
):
    """Returns true when the requirements.txt file was changed"""
    return (stored_requirements_timestamp is None) or (
        current_requirements_timestamp > stored_requirements_timestamp
    )


def is_script_outdated(current_script_timestamp, stored_script_timestamp):
    """Returns true when this script was changed"""
    return (stored_script_timestamp is None) or (
        current_script_timestamp > stored_script_timestamp
    )


def main(arguments):
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument(
        "-d",
        "--venv_root",
        help="Root dir where the environment is generated",
        default=".",
        type=str,
    )

    parser.add_argument(
        "-n",
        "--venv_name",
        help="Name of the environment to be generated",
        default="venv",
        type=str,
    )

    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Activates verbose output",
        default=False,
    )

    parser.add_argument(
        "-c",
        "--clear",
        action="store_true",
        help="Delete the contents of the environment directory before environment creation.",
        default=False,
    )

    args = parser.parse_args(arguments)
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    venv_dir = os.path.join(args.venv_root, args.venv_name)
    logger.info(f"Directory of venv is {venv_dir}")
    venv_exists = os.path.isdir(venv_dir)

    # Get timestamp of this script
    script_timestamp_file = os.path.join(venv_dir, SCRIPT_TIMESTAMP_FILENAME)
    logger.info(f"Path of script timestamp file is {script_timestamp_file}")

    current_script_timestamp = get_file_timestamp(SCRIPT_PATH)
    stored_script_timestamp = (
        float(Path(script_timestamp_file).read_text().strip())
        if os.path.isfile(script_timestamp_file)
        else None
    )

    script_outdated = is_script_outdated(
        current_script_timestamp, stored_script_timestamp
    )
    clear_venv = args.clear or script_outdated
    if (not venv_exists) or clear_venv:
        print("Creating virtual environment...")
        create_virtualenv(venv_dir, clear_venv)

    #----

    # Get timestamp of requirements.txt
    requirements_timestamp_file = os.path.join(
        venv_dir, REQUIREMENTS_TIMESTAMP_FILENAME
    )
    logger.info(f"Path of requirements timestamp file is {requirements_timestamp_file}")

    current_requirements_timestamp = get_file_timestamp(REQUIREMENTS_FILE)
    stored_requirements_timestamp = (
        float(Path(requirements_timestamp_file).read_text().strip())
        if os.path.isfile(requirements_timestamp_file)
        else None
    )

    # Check requirements.txt
    req_txt_outdated = is_requirement_txt_outdated(
        current_requirements_timestamp, stored_requirements_timestamp
    )

    if req_txt_outdated or script_outdated:
        update_virtual_env(
            venv_dir,
            requirements_timestamp_file,
            current_requirements_timestamp,
            script_timestamp_file,
            current_script_timestamp,
        )
    else:
        print("Virtual environment is up-to-date.")

    return 0  # successfull termination


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))