fzf-lua/lua/fzf-lua/shell.lua

-- modified version of:
-- https://github.com/vijaymarupudi/nvim-fzf/blob/master/lua/fzf/actions.lua
local uv = vim.loop
local path = require "fzf-lua.path"
local libuv = require "fzf-lua.libuv"

local M = {}

local _counter = 0
local _registry = {}

function M.register_func(fn)
  _counter = _counter + 1
  _registry[_counter] = fn
  return _counter
end

function M.get_func(counter)
  return _registry[counter]
end

-- creates a new address to listen to messages from actions. This is important,
-- if the user is using a custom fixed $NVIM_LISTEN_ADDRESS. Different neovim
-- instances will then use the same path as the address and it causes a mess,
-- i.e. actions stop working on the old instance. So we create our own (random
-- path) RPC server for this instance if it hasn't been started already.
-- NOT USED ANYMORE, we use `vim.g.fzf_lua_server` instead
-- local action_server_address = nil

function M.raw_async_action(fn, fzf_field_expression, debug)

  if not fzf_field_expression then
    fzf_field_expression = "{+}"
  end

  local receiving_function = function(pipe_path, ...)
    local pipe = uv.new_pipe(false)
    local args = {...}
    uv.pipe_connect(pipe, pipe_path, function(err)
      if err then
        error(string.format("pipe_connect(%s) failed with error: %s",
          pipe_path, err))
      else
        vim.schedule(function ()
          fn(pipe, unpack(args))
        end)
      end
    end)
  end

  local id = M.register_func(receiving_function)

  -- this is for windows WSL and AppImage users, their nvim path isn't just
  -- 'nvim', it can be something else
  local nvim_bin = vim.v.progpath

  local call_args = ("fzf_lua_server=[[%s]], fnc_id=%d %s"):format(
    vim.g.fzf_lua_server, id, debug and ", debug=true" or "")

  local action_cmd = ("%s -n --headless --clean --cmd %s %s"):format(
    libuv.shellescape(nvim_bin),
    libuv.shellescape(("lua loadfile([[%s]])().rpc_nvim_exec_lua({%s})")
      :format(path.join{vim.g.fzf_lua_directory, "shell_helper.lua"}, call_args)),
    fzf_field_expression)

  return action_cmd, id
end

function M.async_action(fn, fzf_field_expression, debug)
  local action_string, id = M.raw_async_action(fn, fzf_field_expression, debug)
  return vim.fn.shellescape(action_string), id
end

function M.raw_action(fn, fzf_field_expression, debug)

  local receiving_function = function(pipe, ...)
    local ret = fn(...)

    local on_complete = function(_)
      -- We are NOT asserting, in case fzf closes
      -- the pipe before we can send the preview
      -- assert(not err)
      uv.close(pipe)
    end

    if type(ret) == "string" then
      uv.write(pipe, ret, on_complete)
    elseif type(ret) == nil then
      on_complete()
    elseif type(ret) == "table" then
      if not vim.tbl_isempty(ret) then
        uv.write(pipe, vim.tbl_map(function(x) return x.."\n" end, ret), on_complete)
      else
        on_complete()
      end
    else
      uv.write(pipe, tostring(ret) .. "\n", on_complete)
    end
  end

  return M.raw_async_action(receiving_function, fzf_field_expression, debug)
end

function M.action(fn, fzf_field_expression ,debug)
  local action_string, id = M.raw_action(fn, fzf_field_expression, debug)
  return vim.fn.shellescape(action_string), id
end

M.preview_action_cmd = function(fn, fzf_field_expression, debug)
  local action_string, id =
    M.raw_preview_action_cmd(fn, fzf_field_expression, debug)
  return vim.fn.shellescape(action_string), id
end

M.raw_preview_action_cmd = function(fn, fzf_field_expression, debug)

  return M.raw_async_action(function(pipe, ...)

    local function on_finish(_, _)
      if pipe and not uv.is_closing(pipe) then
        uv.close(pipe)
        pipe = nil
      end
    end

    local function on_write(data, cb)
      if not pipe then
        cb(true)
      else
        uv.write(pipe, data, cb)
      end
    end

    return libuv.spawn({
        cmd = fn(...),
        cb_finish = on_finish,
        cb_write = on_write,
      }, false)

  end, fzf_field_expression, debug)
end

M.reload_action_cmd = function(opts, fzf_field_expression)

  if opts.fn_preprocess and type(opts.fn_preprocess) == 'function' then
    -- run the preprocessing fn
    opts = vim.tbl_deep_extend("keep", opts, opts.fn_preprocess(opts))
  end

  return M.raw_async_action(function(pipe, args)

    -- get the type of contents from the caller
    local reload_contents = opts.__fn_reload(args[1])
    local write_cb_count = 0
    local pipe_want_close = false

    -- local on_finish = function(code, sig, from, pid)
      -- print("finish", pipe, pipe_want_close, code, sig, from, pid)
    local on_finish = function(_, _, _, _)
      if not pipe then return end
      pipe_want_close = true
      if write_cb_count==0 then
        -- only close if all our uv.write calls are completed
        uv.close(pipe)
        pipe = nil
      end
    end

    local on_write = function(data, cb, co)
      -- pipe can be nil when using a shell command with spawn
      -- and typing quickly, the process will terminate and
      assert(not co or (co and pipe and not uv.is_closing(pipe)))
      if not pipe then return end
      if not data then
        on_finish(nil, nil, 5)
        if cb then cb(nil) end
      else
        write_cb_count = write_cb_count + 1
        uv.write(pipe, tostring(data), function(err)
          write_cb_count = write_cb_count - 1
          if co then coroutine.resume(co) end
          if cb then cb(err) end
          if err then
            -- force close on error
            write_cb_count = 0
            on_finish(nil, nil, 2)
          end
          if write_cb_count == 0 and pipe_want_close then
            on_finish(nil, nil, 3)
          end
        end)
        -- yield returns when uv.write compeletes
        -- or when a new coroutine calls resume(1)
        if co and coroutine.yield() == 1 then
          -- we have a new routine in opts.__co, this
          -- routine is no longer relevant so kill it
          write_cb_count = 0
          on_finish(nil, nil, 4)
        end
      end
    end

    if type(reload_contents) == 'string' then
      -- string will be used as a shell command
      -- terminate previously running commands
      libuv.process_kill(opts.__pid)
      opts.__pid = nil

      -- spawn/async_spawn already async, no need to send opts.__co
      -- also, we can't call coroutine.yield inside a libuv callback
      -- due to: "attempt to yield across C-call boundary"
      libuv.async_spawn({
        cwd = opts.cwd,
        cmd = reload_contents,
        cb_finish = on_finish,
        cb_write = on_write,
        cb_pid = function(pid) opts.__pid = pid end,
        -- must send false, 'coroutinify' adds callback as last argument
        -- which will conflict with the 'fn_transform' argument
      }, opts.__fn_transform or false)
    else
      -- table or function runs in a coroutine
      -- which isn't required for 'libuv.spawn'
      coroutine.wrap(function()

        if opts.__co then
          local costatus = coroutine.status(opts.__co)
          if costatus ~= 'dead' then
            -- the previous routine is either 'running' or 'suspended'
            -- return 1 from yield to signal abort to 'on_write'
            coroutine.resume(opts.__co, 1)
          end
          assert(coroutine.status(opts.__co) == 'dead')
        end
        -- reset var to current running routine
        opts.__co = coroutine.running()

        -- callback with newline
        local on_write_nl = function(data, cb)
          data = data and tostring(data) .. "\n" or nil
          return on_write(data, cb)
        end

        -- callback with newline and coroutine
        local on_write_nl_co = function(data, cb)
          data = data and tostring(data) .. "\n" or nil
          return on_write(data, cb, opts.__co)
        end

        -- callback with coroutine (no NL)
        local on_write_co = function(data, cb)
          return on_write(data, cb, opts.__co)
        end


        if type(reload_contents) == 'table' then
          for _, l in ipairs(reload_contents) do
            on_write_nl_co(l)
          end
          on_finish()
        elseif type(reload_contents) == 'function' then
          -- by default we use the async callbacks
          if opts.func_async_callback ~= false then
            reload_contents(on_write_nl_co, on_write_co)
          else
            reload_contents(on_write_nl, on_write)
          end
        else
        end

      end)()
    end

  end, fzf_field_expression, opts.debug)
end

return M