From c76e25df6573a32fe5764693adbc572c8c89f620 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=B3n=C3=A1n=20Carrigan?= <rcarriga@tcd.ie>
Date: Sun, 21 Aug 2022 21:28:16 +0100
Subject: [PATCH] feat: custom python command

See #14
---
 README.md                   |  7 ++++++-
 lua/neotest-python/base.lua |  3 ---
 lua/neotest-python/init.lua | 31 ++++++++++++++++++++++++++++---
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 8623672..eec1dfe 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,12 @@ require("neotest").setup({
         -- Runner to use. Will use pytest if available by default.
         -- Can be a function to return dynamic value.
         runner = "pytest",
-
+        -- Custom python path for the runner.
+        -- Can be a string or a list of strings.
+        -- Can also be a function to return dynamic value.
+        -- If not provided, the path will be inferred by checking for 
+        -- virtual envs in the local directory and for Pipenev/Poetry configs
+        python = ".venv/bin/python"
         -- Returns if a given file path is a test file.
         -- NB: This function is called a lot so don't perform any heavy tasks within it.
         is_test_file = function(file_path)
diff --git a/lua/neotest-python/base.lua b/lua/neotest-python/base.lua
index 600688a..d7e619a 100644
--- a/lua/neotest-python/base.lua
+++ b/lua/neotest-python/base.lua
@@ -25,9 +25,6 @@ local python_command_mem = {}
 
 ---@return string[]
 function M.get_python_command(root)
-  if not root then
-    root = vim.loop.cwd()
-  end
   if python_command_mem[root] then
     return python_command_mem[root]
   end
diff --git a/lua/neotest-python/init.lua b/lua/neotest-python/init.lua
index 918a54d..e9d5ab4 100644
--- a/lua/neotest-python/init.lua
+++ b/lua/neotest-python/init.lua
@@ -32,6 +32,13 @@ local function get_strategy_config(strategy, python, program, args)
   end
 end
 
+local get_python = function(root)
+  if not root then
+    root = vim.loop.cwd()
+  end
+  return base.get_python_command(root)
+end
+
 local get_args = function()
   return {}
 end
@@ -79,7 +86,7 @@ function PythonNeotestAdapter.discover_positions(path)
      @namespace.definition
   ]]
   local root = PythonNeotestAdapter.root(path)
-  local python = base.get_python_command(root)
+  local python = get_python(root)
   local runner = get_runner(python)
   return lib.treesitter.parse_positions(path, query, {
     require_namespaces = runner == "unittest",
@@ -98,7 +105,7 @@ function PythonNeotestAdapter.build_spec(args)
   x:close()
 
   local root = PythonNeotestAdapter.root(position.path)
-  local python = base.get_python_command(root)
+  local python = get_python(root)
   local runner = get_runner(python)
   local stream_data, stop_stream = lib.files.stream_lines(stream_path)
   local script_args = vim.tbl_flatten({
@@ -167,7 +174,25 @@ end
 setmetatable(PythonNeotestAdapter, {
   __call = function(_, opts)
     is_test_file = opts.is_test_file or is_test_file
-    if type(opts.args) == "function" or (type(opts.args) == "table" and opts.args.__call) then
+    if opts.python then
+      get_python = function(root)
+        local python = opts.python
+
+        if is_callable(opts.python) then
+          python = opts.python(root)
+        end
+
+        if type(python) == "string" then
+          return { python }
+        end
+        if type(python) == "table" then
+          return python
+        end
+
+        return base.get_python(root)
+      end
+    end
+    if is_callable(opts.args) then
       get_args = opts.args
     elseif opts.args then
       get_args = function()