Context

I recently installed NixOS on my machine. I am currently working on some Python project using Visual Studio Code and was previously relying on Python virtual environments and a requirements.txt file to manage my dependencies. Visual Studio Code was smart enough to probe the working directory for virtual environments, and automatically use the one found as the default interpreter. However, using virtual environments and pip goes against Nix’s philosophy which provide a more powerful tool for development environments: nix-shell. The first part of this post presents how to setup this new environment, the second part digs a bit deeper into the internals of Python environments created by Nix.

Creating a Python Environment with nix-shell

As recommended when working on a Python project, I am using nix-shell to create a Python environment through the following Nix expression:

# in shell.nix
with import <nixpkgs> { };
let
  pythonEnv = python311.withPackages (ps: [
    ps.pandas
    ps.requests
    ps.gensim
    ps.numpy
    ps.matplotlib
    ps.mysql-connector
    ps.scikit-learn
    ps.sqlparse
    ps.tqdm
    ps.scipy
    ps.torch
    ps.transformers
    ps.sentence-transformers
    ps.torch-geometric
    ps.networkx
    ps.matplotlib-venn
    ps.jupyter
    ps.notebook
    ps.ipython
  ]);
in
pkgs.mkShell rec {
  packages = [
    pythonEnv
  ];
}

withPackages is a simplified wrapper around python.buildEnv that allows to create a Python environment by including the specified packages (in this case, pandas, request, gensim, …). Running nix-shell shell.nix creates an environment (through the resulting nix-shell) with access to a custom Python executable that includes our dependencies.

nix-shell is just another mechanism for dependency management that allows to share a developing environment across machines (a reproducible one when coupled with version pinning). A major benefit of using nix-shell is the ability to add dependencies that are not Python packages but still necessary for the project. For instance, we could add the graphviz package to manipulate created graph by networkx to the package field. The graphviz application would be included in our developpment environment but not globally.

Visual Studio Code integration

Now, when opening an instance of Visual Studio Code outside of our nix-shell, Visual Studio Code only sees the global Python interpreter without any dependencies: imports are broken so as semantic highlighting, and code navigation.

Visual Studio Code does not have access to the Python environment built with nix-shell. This means that to allow Visual Studio Code to access our environment we must invoke it from the nix-shell. Only then, it is able to find the Python environment, and we can manually select from the interpreter selection box.

However, Visual Studio Code sometimes keep in memory previously used interpreters. When adding a new package as a dependency to our project, a new Python environment is created and the interpreter path is different. Quickly, Visual Studio Code will reference many Python environments and knowing which one is is the correct one is not always trivial. Interpreter cache can be flushed using commands (Python: clear workspace interpreter settings & Python: Clear cache and reload windows) but that requires many inputs, and we still have to manually select the correct interpreter.

Visual Studio Code automatically choose a python interpreter by probing (in that order): virtual environments in the workspace, workspace related virtual environments and finally, globally installed interpreters. I don’t want to use virtual environments as Nix already built a Python environment with all necessary packages, nor can I rely on the automatic selection of the correct interpreter amongst the global ones, I would rather specify directly to Visual Studio Code which interpreter to use. Luckily, a Visual Studio Code setting does exactly that: python.defaultInterpreterPath. In your workspace settings you can create (or modify) a .vscode/settings.json file where you set the path to the default interpreter to some environment variable:

{
    "python.defaultInterpreterPath" : "${CUSTOM_INTERPRETER_PATH}"
}

And add a hook to the shell.nix file to export the environment variable accordingly:

[...]

pkgs.mkShell rec {
  packages = [
    pythonEnv
  ];
  
  shellHook = ''
    export CUSTOM_INTERPRETER_PATH="${pythonEnv}/bin/python"
  ''; 
}

These changes force Visual Studio Code into selecting the Python environment created by Nix for the workspace by using the CUSTOM_INTERPRETER_PATH environment variable. When rebuilding a new Python environment by adding or removing a package, simply re-invoke a new Visual Studio Code instance in the updated nix-shell.

Nix’s Python environment

Now, more on the Python environment. Using nix-shell on the previous Nix expression creates a temporary environment with access to all defined dependencies. Through withPackages, we are creating a custom Python executable that has access to the declared packages: a Python environment.

[user@~/repos/some-project]$ nix-shell shell.nix

[nix-shell:~/repos/some-project]$ which python3 

/nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/bin/python3

$ python3 
Python 3.11.10 (main, Sep  7 2024, 01:03:31) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import transformers 
>>> import torch 
>>> import matplotlib_venn 
>>> 

/nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/bin/python3 is our custom Python executable specific to the declared environment and as we can see, it has access to our dependency packages. If we add or remove a package, a new Python environment will be created (reusing already downloaded or built packages, building or downloading the added ones).

But how does the custom Python3 executable locate the dependencies ? Usually, one can modify the PYTHONPATH environment variable to point towards packages. However, when creating child processes for a different project from the parent shell, these processes inherit the parent’s PYTHONPATH. This inheritance can cause conflicts in the development environment, as packages from the parent’s path may take precedence over the intended versions in the child environment.

Package resolution mechanism depends wether dependencies are included in the Python environment or outside it.

Packages included in Python Environments

Dependencies added to a Python environment through withPackages are first added to the Nix store, and then a symbolic link is added under the Python environment’s site-packages pointing to the module. In our case, the Python environment has been created under /nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/. We can find the locations Python is looking at when search for a package by inspecting sys.path:

[nix-shell:~/repos/some-project]$ python3 
Python 3.11.10 (main, Sep  7 2024, 01:03:31) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sys
>>> print("\n".join(sys.path))
/nix/store/nmqxyr00in2arwrq5qd1qipsanz1yrn5-python3-3.11.10/lib/python311.zip
/nix/store/nmqxyr00in2arwrq5qd1qipsanz1yrn5-python3-3.11.10/lib/python3.11
/nix/store/nmqxyr00in2arwrq5qd1qipsanz1yrn5-python3-3.11.10/lib/python3.11/lib-dynload
/nix/store/nmqxyr00in2arwrq5qd1qipsanz1yrn5-python3-3.11.10/lib/python3.11/site-packages
/nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/lib/python3.11/site-packages
>>> 

Examining the Python environment’s site-packages folder, we can retrieve the symbolic links pointing to all of our dependencies declared in the Python environment (and their runtime dependencies as well):

[nix-shell:~/repos/some-project]$  ls -la /nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/lib/python3.11/site-packages
...
lrwxrwxrwx 1 root root   103 janv.  1  1970 pandas -> 
  /nix/store/ykj2zgln73lsm8wr797wsc8b2mpdp6k8-python3.11-pandas-2.2.3/lib/python3.11/site-packages/pandas
lrwxrwxrwx 1 root root   119 janv.  1  1970 pandas-2.2.3.dist-info -> 
  /nix/store/ykj2zgln73lsm8wr797wsc8b2mpdp6k8-python3.11-pandas-2.2.3/lib/python3.11/site-packages/pandas-2.2.3.dist-info
...

This also reveals the usage of symbolic links by Nix to include packages from the store under our Python environment’s site-package directory. The transparent path resolution can be illustrated by revealing the module’s filepath:

[nix-shell:~/repos/xp-gaur]$ python3 
Python 3.11.10 (main, Sep  7 2024, 01:03:31) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pandas
p>>> print(pandas.__file__)
/nix/store/589qgvh8zrvyda825id44p77nbq8gqpf-python3-3.11.10-env/lib/python3.11/site-packages/pandas/__init__.py
>>> 

Packages declared outside the Python Environment

On the other hand, other Python package dependencies that are not included in the Python environment (for instance, local packages added through callPackage) are imported using sitecustomize. Nix places such modules in the Python environment’s site-packages. During startup, Python tries to import this module to add custom locations in sys.path. In Nix’s implementation, the module reads the NIX_PYTHONPATH environment variable to append local packages (and their dependencies) to sys.path.

Conclusion

This post detailed how to create and manage Python development environments using nix-shell as well as the package resolution logic when declaring dependencies. Through nix-shell, Nix offers a more powerful alternative to traditional Python virtual environments by creating reproducible environments that manage both Python libraries and system dependencies in a declarative way. Moreover, nix-shell is not specific to Python projects but can be used with many languages and frameworks.