modules: add python options to specify the python environment

modules/output.nix

@@ -1,6 +1,7 @@
 {
   lib,
   config,
+  pkgs,
   ...
 }:
 with lib; let
@@ -38,6 +39,24 @@ in {
       description = "Extra packages to be made available to neovim";
     };
 
+    python = {
+      package = mkOption {
+        type = types.package;
+        default = pkgs.python3;
+        description = "";
+      };
+
+      extraPythonPackages = mkOption {
+        type = with types; functionTo (listOf package);
+        default = p: [];
+        defaultText = literalExpression "p: with p; [ ]";
+        description = "Python packages to add to the `PYTHONPATH` of neovim.";
+        example = lib.literalExpression ''
+          p: [ p.numpy ]
+        '';
+      };
+    };
+
     extraConfigLua = mkOption {
       type = types.lines;
       default = "";

tests/test-sources/examples.nix

@@ -1,6 +1,16 @@
 {pkgs, ...}: {
   plain = {};
 
+  python-packages = {
+    python = {
+      package = pkgs.python3;
+      extraPythonPackages = p:
+        with p; [
+          numpy
+        ];
+    };
+  };
+
   simple-plugin = {
     extraPlugins = [pkgs.vimPlugins.vim-surround];
   };

wrappers/modules/output.nix

@@ -103,6 +103,13 @@ in {
         inherit (config) viAlias vimAlias;
         # inherit customRC;
         plugins = normalizedPlugins;
+
+        # Python 3 environment
+        python3Env = let
+          python = config.python.package;
+          inherit (config.python) extraPythonPackages;
+        in
+          python.withPackages extraPythonPackages;
       }
       # Necessary to make sure the runtime path is set properly in NixOS 22.05,
       # or more generally before the commit: