Python Project Template

{
  # Arguments that can be modified by the caller
  system ? builtins.currentSystem,
  # We specify which version of nixpkgs to use as input. This is a specific commit that we can
  # view on GitHub.
  # The revision specified here should be the same as in the flake.lock file generated by the
  # flake.nix file below for the two shells to be identical.
  nixpkgs ? fetchTarball
     # This is the most recent commit on the 23.11 channel as I write this.
     "https://github.com/NixOS/nixpkgs/archive/057f9aecfb71c4437d2b27d3323df7f93c010b7e.tar.gz"
}:

let
  # This pkgs section could also be put into the inputs above,
  # so that the user would be able to also change pkgs.
  pkgs = import nixpkgs {
    inherit system;

    # Unlike with the flake version below, it is important that
    # that we set both config and overlays to maintain purity.
    # If we don't, they try and load the users configs and overlays.
    config = { };

    overlays = [
      # To change a package version, we add an overlay that calls first
      # the override method on the wrapped Python interpreter that we are
      # using, as we want to override the packages that it provides. Then,
      # we add a packageOverride to modify scikit-learn.
      (final: prev: {
        python311 = prev.python311.override {
          packageOverrides = pyfinal: pyprev: {
            "scikit-learn" = pyprev.scikit-learn.overridePythonAttrs (old: rec {
              pname = "scikit-learn";
              version = "1.4.0";

              # Here we fetch a more recent version from Pypi servers.
              # We could also pull a version from Github.
              src = pyprev.fetchPypi {
                inherit pname version;

                # If the hash is unknown, we can leave it blank, add nix
                # will tell us the hash.
                hash = "sha256-1Dc8mE66IOOTIW7dUaPj7t5Wy+k9QkdRbSBWQ8O5MSE=";
              };
            });
          };
        };
      })

      # We could add additional overlays here...
    ];

  };

in
pkgs.mkShellNoCC {

  buildInputs = with pkgs; [
    # We tell the shell to include our chosen Python interpreter
    # and all the packages we require
    (python311.withPackages (ps: [
      ps.pandas
      ps.numpy
      ps.scikit-learn # This is now version 1.4.0
    ]))
  ];

}

$> nix-shell --pure
$> python
Python 3.11.6 (main, Oct  2 2023, 13:45:54) [GCC 12.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sklearn
>>> sklearn.__version__
'1.4.0'
>>> import pandas as pd
>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'torch'

{
  description = "A modified Python wrapper";

  inputs = {
    # Specify which channel to use
    # The exact commit will then be saved in the flake.lock file.
    nixpkgs.url = "github:NixOS/nixpkgs/23.11";
    utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, ... }@inputs: inputs.utils.lib.eachDefaultSystem (system: let
    pkgs = import nixpkgs {

      inherit system;

      overlays = [

        # To change a package version, we add an overlay that calls first
        # the override method on the wrapped Python interpreter that we are
        # using, as we want to override the packages that it provides. Then,
        # we add a packageOverride to modify scikit-learn.
        (final: prev: {
          python311 = prev.python311.override {
            packageOverrides = pyfinal: pyprev: {
              "scikit-learn" = pyprev.scikit-learn.overridePythonAttrs (old: rec {
                pname = "scikit-learn";
                version = "1.4.0";

                # Here we fetch a more recent version from Pypi servers.
                # We could also pull a version from Github.
                src = pyprev.fetchPypi {
                  inherit pname version;

                  # If the hash is unknown, we can leave it blank, add nix
                  # will tell us the hash.
                  hash = "sha256-1Dc8mE66IOOTIW7dUaPj7t5Wy+k9QkdRbSBWQ8O5MSE=";
                };
              });
            };
          };
        })

        # We could add additional overlays here...

      ];

    };
  in {

    # Finally, we set up our development shell
    devShells.default = pkgs.mkShellNoCC {

      packages = with pkgs; [
        # We tell the shell to include our chosen Python interpreter
        # and all the packages we require
        (python311.withPackages (ps: [
          ps.pandas
          ps.numpy
          ps.scikit-learn # This is now version 1.4.0
        ]))
      ];

    };

  });
}

$> nix develop
$> python
Python 3.11.6 (main, Oct  2 2023, 13:45:54) [GCC 12.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sklearn
>>> sklearn.__version__
'1.4.0'
>>> import pandas as pd
>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'torch'