Sep 8, 2022

Fixing Python Import Resolution in Nix with Direnv

Introduction

I've started using development environments with nix-shell for my personal projects. I greatly prefer the consistency of dropping into a nix-shell, over reading Python's venv manual for the umpteenth time, then yak-shaving into whether I should be using pyenv-virtualenv instead.

Here's a typical shell.nix file for Python 3.10 and some common packages like numpy.

{ pkgs ? import <nixpkgs> { } }:

let
  my-python = pkgs.python310;
  python-with-my-packages = my-python.withPackages
    (p: with p; [ numpy pytorch matplotlib requests python-dotenv ]);
in pkgs.mkShell {
  buildInputs = [
    python-with-my-packages
  ];
  shellHook = ''
    PYTHONPATH=${python-with-my-packages}/${python-with-my-packages.sitePackages}
  '';
}

Relying on this, it's likely you will see your editor (specifically, your LSP implementation of choice) highlighting import errors, like so1:

/img/python-import-errors.png

However, there aren't actually any errors at runtime. When run inside the nix-shell environment, we certainly have pytorch, matplotlib, etc. So, how do we make our LSP server aware of the packages we have in our nix-shell environment?

Solution

TLDR:

  1. Install direnv.
  2. Install an editor extension/plugin for direnv, e.g. emacs-direnv. If you use Doom Emacs, simply enable the direnv module.
  3. echo use_nix > .envrc in your project directory.
  4. direnv allow . in your project directory.

Now reload your editor.

Explanation

Pyright, my language server of choice for Python, gets the installed packages directory from the PYTHONPATH environment variable. So, a subpar approach would be to hard-code that to the location in the /nix/store that contains the packages. Anytime you changed the Python version, for example, this value would change. Instead of this, we can use direnv's clever integration with Nix. The described use case is automatically loading environment variables in a shell; our use case is automatically loading environment variables into our editor.

So every time you open your project, direnv sees use_nix in .envrc, resolves your shell.nix file, and injects any environment variables (PYTHONPATH) into your editor, using your direnv editor plugin.

/img/python-import-success.png

That's it! Pyright sees PYTHONPATH and we have working import resolution.

Follow-up

  • nix-direnv may be quicker than direnv, reducing the time to resolve the nix-shell environment.
  • lorri appears to be a more feature-rich replacement for the virtual environment use case.
1

In this case, the packages that aren't highlighted are the ones previously installed on my machine.