DIY CLI-tool: Find and Activate Virtual Environments

In this post, I will cover my most used CLI-tool. It finds a virtual environment, created using virtualenv, in your current directory and activates it. It also allows you to pick which environment to activate if multiple where detected. Although this doesn't really save much time, it gives a much nicer working experience when jumping between many different projects. I will be talking about these tools from the perspective of my setup, described in parts one and two of this series. The target of this script is Windows, however, I would imagine it is easily ported to Linux as well.

Here is an example of running the tool with the recursive flag, finding multiple environments: Alt Text

If you want all the code, I put it up on GitHub in a repo called activate_env. At the time of writing the code in the repo is the same as the one shown in this post. However, it might have changed when you are reading this. In the end, the core design will most likely be the same even if some implementation details might change.

The complexity of this tool comes from the needed interaction with your current shell. This, as far as I know, is not possible to do from Python. Therefore, it has three files in three different languages. The main file is:

_find_virtual_env.py, which does almost all the work

Two two shell files for wrapping the python script and then activates the environment:

activate_env.bat, activates the environment when running cmd-shell
activate_env.ps1, activates the environment when running powershell

I will take one file at the time starting with the meat of the tool, _find_virtual_env.py.

`_find_virtual_env.py`

This parses command-line arguments, find all virtual environments in the directory and allow the user to pick which environment to activate. As this script will be wrapped by two other shell-scripts, I decided to have all interactions with the user go through stderr. Then it will just write the path to the activation script of the environment to stdout which is picked up by the wrappers. I will go through the file from top to bottom, one function at the time and then finally the dunder main, __main__.

As you might expect, we start off with the imports. As you can see, this tool only depends on the standard library. Well, except for you having virtual environments, but I guess that goes without saying.

import sys
from pathlib import Path
from contextlib import contextmanager
from typing import List
from argparse import ArgumentParser
import signal

As I stated earlier, I want to write almost all output to stderr. Therefore, I created a context manager to help me with this. This re-directs everything written to stdout to stderr while in the context-block and returns everything to normal when exiting.

@contextmanager
def stderr_as_out():
    stdout = sys.stdout
    sys.stdout = sys.stderr
    yield
    sys.stdout = stdout

Next, I have a function which we will later register to react to the SIGINT signal (pressing ctrl-c). This was added to clean up the output when exiting from the program using an interrupt. It prints a newline to stderr and returns a failing status code. Maybe, the code actually should be zero in this case? However, I am currently using success (0) only when it actually produced an environment to be activated.

def handle_sigint(*args):
    """Makes output cleaner sent SIGINT (ctrl-c)"""
    with stderr_as_out():
        print("")
    sys.exit(1)

The find_envs function searches for virtual environments in the supplied directory. As the environment can be named whatever you want, it looks for a folder containing the path Scripts/activate. So, if you had an environment named env this would be env/Scripts/activate. This is the folder structure of virtual environments, at least on Windows. It uses glob for finding all paths matching this structure, by default only at the current directory level. Because it was easy to add, it also supports the use of recursive search. This just switches the glob-expression from one leading start to two.

def find_envs(cur_dir: Path, recursive: bool = False) -> List[Path]:
    search_pattern = "**/Scripts/activate" if recursive else "*/Scripts/activate"
    return [env_path for env_path in cur_dir.glob(search_pattern)]

If the program finds multiple environments I want the user to choose which to activate. For this I have two functions, the first one prints the choices with some added styling. Each choice has a index (used to make the choice), the name of the virtual environment and path to the environment.

def print_envs(envs: List[Path]) -> None:
    max_name_len = max([len(env.parent.parent.name) for env in envs])
    for idx, env in enumerate(envs):
        env_dir = env.parent.parent
        name_str = f"({env_dir.name})"
        print(f"[{idx}]: {name_str:^{max_name_len+3}} - {env_dir}")

The next function accept the user input, repeating the question prompt until a valid index is supplied. It then returns the environment path corresponding to the index supplied by the user.

def ask_user_to_choose_env(envs: List[Path]) -> Path:
    max_idx = len(envs) - 1
    while True:
        idx = input("Activate environment (index): ")
        try:
            idx = int(idx)
        except ValueError:
            print("Index must be a integer")
            continue

        if (idx < 0) or (idx > max_idx):
            print(f"Index must be in range [0, {max_idx}]")
            continue
        else:
            break

    return envs[idx]

Now to the main, where all the pieces are put together. First, we register the SIGINT handler show earlier. Secondly, the argument parser is created, which automatically adds the -h flag for help messages. Then a -r flag is added, which is a Boolean flag used to trigger the recursive version of find_envs.

The majority of the program is then run as a part of stderr_as_out-context block. Parsing of the arguments is also done here, as the -h flag otherwise would trigger some wired output from the wrapping scripts. The find_env function runs and depending on the result the behavior differs. If no environment was found we print this information and exit the program. If only a single one was found it is activated and otherwise, we ask the user to pick an environment to activate. The final step is then to write this to stdout using the final print-statement, which is outside the context-block.

if __name__ == "__main__":
    signal.signal(signal.SIGINT, handle_sigint)
    parser = ArgumentParser(
        prog="activate",
        description="""
            Activates a virtual environment if found in the current directory. 
            If multiple are found the user may choose the environment. 
            All user-communication is written through stderr to not conflict with calling programs.
        """.strip(),
    )
    parser.add_argument(
        "-r",
        "--recursive",
        action="store_true",
        default=False,
        help="Recursivly searches through the directory for virtual environments",
    )
    with stderr_as_out():
        # We re-direct the parsing output to stderr as e.g -h argument will produce ouput
        # This would break the ps1/bat files which uses this script
        args = parser.parse_args()

        envs = find_envs(cur_dir=Path("."), recursive=args.recursive)

        if not envs:
            print("No virtual envs found")
            exit(1)
        elif len(envs) > 1:
            print_envs(envs)
            env = ask_user_to_choose_env(envs)
        else:
            env = envs[0]

    # Write env to stdout for the wrapping shell scripts to grab it
    print(str(env))
    exit(0)

`activate_envs.ps1`

Next up, is the powershell script used to actually activate the chosen environment, when using powershell. This is a rather straight forward script. It runs the program, passing through all arguments to it and stores everything it has written to stdout in a variable. As we know the script only writes the final environment choice to stdout it is not much more work to be done. We check if it contains anything and if it does the Invoke-Expression statement will run the activate script. This is exactly what you would have done when manually activating the environment.

$env = _find_virtual_env.py $args
if ($env){
    Invoke-Expression $env
}

`activate_envs.bat`

Now, this file was the hardest one to write. I have to say batch is not a pretty language, and I hope you like GOTO! This time I will actually show you the code first and then try to explain it, so without further ado:

@echo off

FOR /F "tokens=*" %%F IN ('_find_virtual_env.py %*') DO (
    SET __choosen_virtual_env__=%%F
)

if [%__choosen_virtual_env__%]==[] GOTO EXIT 

call %__choosen_virtual_env__%


:EXIT
SET __choosen_virtual_env__=

The first step is to turn off echo, which otherwise will print every single line of your script to the terminal as it is running. Now I will decompose the for-loop.

The /F flag, I think stands for file and you treat stdout as reading from a file.
The parameter "tokens=*", tells it to read all tokens from the file.
%%F of course declares the loop variable, which must be a single letter.
'_find_virtual_env.py %*', runs the python script, passing all parameters to it
The SET-statement stores the result in a new variable named __choosen_virtual_env__

I choose this rather complicated variable name as I did not want it to clash with any other global variable, because, well everything is global. There is a method for creating a local scope with startlocal and endlocal. However, this does not work when trying to activate the environment. This most likely comes from the environment activation depending on setting multiple environment variables. These, of course, are also affected by the local-scope created in this file and are destroyed when the scope ends. So, a global variable it is.

The if-statement checks if the variable is empty and if it is, jumps to the EXIT label. If it wasn't empty we will reach the call-statement which runs the activation script, again similar to how you would have done it manually. The final step, after the EXIT label, is to set the variable to nothing. This very pretty syntax basically deletes the variable.

Summary

We use Python for the heavy lifting, searching for virtual environments using glob and allowing the user to choose which environment to activate if multiple were detected. All communication from the python script to the user is done using stderr and it only writes the final path to stdout. Then we wrap it in thin wrappers for each shell, cmd and powershell. The wrapper just takes a single path from the Python script and runs it, activating the virtual environment.