From 59e1e8637c0bb0945dc73c9d73b6d395054baccd Mon Sep 17 00:00:00 2001
From: Smaug123 <patrick+github@patrickstevens.co.uk>
Date: Fri, 29 Mar 2024 11:58:52 +0000
Subject: [PATCH] More python stuff

---
 home-manager/home.nix               |  2 +-
 home-manager/nvim/python.lua        | 69 +++++++++++++++++++++++++++++
 home-manager/nvim/venv-selector.lua |  1 +
 3 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 home-manager/nvim/python.lua

diff --git a/home-manager/home.nix b/home-manager/home.nix
index a09a38f..a9fca46 100644
--- a/home-manager/home.nix
+++ b/home-manager/home.nix
@@ -272,7 +272,7 @@
     vimdiffAlias = true;
     withPython3 = true;
 
-    extraLuaConfig = builtins.readFile ./nvim/build-utils.lua + "\n" + builtins.readFile ./nvim/dotnet.lua + "\n" + builtins.replaceStrings ["%PYTHONENV%"] ["${pythonEnv}"] (builtins.readFile ./nvim/init.lua);
+    extraLuaConfig = builtins.readFile ./nvim/build-utils.lua + "\n" + builtins.readFile ./nvim/dotnet.lua + "\n" + builtins.replaceStrings ["%PYTHONENV%"] ["${pythonEnv}"] (builtins.readFile ./nvim/init.lua) + "\n" + builtins.readFile ./nvim/python.lua;
 
     package = nixpkgs.neovim-nightly;
   };
diff --git a/home-manager/nvim/python.lua b/home-manager/nvim/python.lua
new file mode 100644
index 0000000..997f5de
--- /dev/null
+++ b/home-manager/nvim/python.lua
@@ -0,0 +1,69 @@
+local function pytest_on_line(_, _, _) end
+local function pytest_on_complete(_, code, _)
+	if code ~= 0 then
+		print("Exit code " .. code)
+	end
+end
+
+function RunPythonTestAtCursor()
+	local api = vim.api
+
+	-- Get the current buffer and cursor position
+	local bufnr = api.nvim_get_current_buf()
+	local line_nr = api.nvim_win_get_cursor(0)[1]
+	local filename = api.nvim_buf_get_name(bufnr)
+
+	-- Read the file content
+	local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
+
+	-- Find the test function
+	local test_name = nil
+	for i = line_nr, 1, -1 do
+		local line = lines[i]
+		if line:match("^def test_") then
+			test_name = line:match("^def (%S+)%(")
+			break
+		end
+	end
+
+	if test_name then
+		-- Run pytest for the found test function
+		local context = BuildUtils.create_window()
+		BuildUtils.run(
+			"pytest",
+			{ filename .. "::" .. test_name },
+			"Run PyTest (" .. test_name .. ")",
+			context,
+			pytest_on_line,
+			pytest_on_complete
+		)
+	else
+		print("No test function found at or above line " .. line_nr)
+	end
+end
+
+function RunPythonTestsInFile()
+	local file_path = vim.fn.expand("%:p")
+	local context = BuildUtils.create_window()
+	BuildUtils.run("pytest", { file_path }, "Run PyTest", context, pytest_on_line, pytest_on_complete)
+end
+
+function RunAllPythonTests()
+	local context = BuildUtils.create_window()
+	BuildUtils.run("pytest", {}, "Run PyTest", context, pytest_on_line, pytest_on_complete)
+end
+
+do
+	local whichkey = require("which-key")
+	whichkey.register({
+		p = {
+			name = "Python-related commands",
+			t = {
+				"Run Python tests",
+				f = { RunPythonTestsInFile, "Run Python tests in the current file" },
+				a = { RunAllPythonTests, "Run all Python tests" },
+				c = { RunPythonTestAtCursor, "Run the Python test under the cursor" },
+			},
+		},
+	}, { prefix = vim.api.nvim_get_var("maplocalleader") })
+end
diff --git a/home-manager/nvim/venv-selector.lua b/home-manager/nvim/venv-selector.lua
index 6b76468..6f13291 100644
--- a/home-manager/nvim/venv-selector.lua
+++ b/home-manager/nvim/venv-selector.lua
@@ -3,6 +3,7 @@ local venv_selector = require("venv-selector")
 venv_selector.setup({
 	changed_venv_hooks = { venv_selector.hooks.pyright },
 	name = { "venv", ".venv" },
+	search_venv_managers = true,
 })
 
 vim.api.nvim_create_autocmd("VimEnter", {