_system.lua 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393
  1. local uv = vim.uv
  2. --- @class vim.SystemOpts
  3. --- @field stdin? string|string[]|true
  4. --- @field stdout? fun(err:string?, data: string?)|false
  5. --- @field stderr? fun(err:string?, data: string?)|false
  6. --- @field cwd? string
  7. --- @field env? table<string,string|number>
  8. --- @field clear_env? boolean
  9. --- @field text? boolean
  10. --- @field timeout? integer Timeout in ms
  11. --- @field detach? boolean
  12. --- @class vim.SystemCompleted
  13. --- @field code integer
  14. --- @field signal integer
  15. --- @field stdout? string
  16. --- @field stderr? string
  17. --- @class vim.SystemState
  18. --- @field cmd string[]
  19. --- @field handle? uv.uv_process_t
  20. --- @field timer? uv.uv_timer_t
  21. --- @field pid? integer
  22. --- @field timeout? integer
  23. --- @field done? boolean|'timeout'
  24. --- @field stdin? uv.uv_stream_t
  25. --- @field stdout? uv.uv_stream_t
  26. --- @field stderr? uv.uv_stream_t
  27. --- @field stdout_data? string[]
  28. --- @field stderr_data? string[]
  29. --- @field result? vim.SystemCompleted
  30. --- @enum vim.SystemSig
  31. local SIG = {
  32. HUP = 1, -- Hangup
  33. INT = 2, -- Interrupt from keyboard
  34. KILL = 9, -- Kill signal
  35. TERM = 15, -- Termination signal
  36. -- STOP = 17,19,23 -- Stop the process
  37. }
  38. ---@param handle uv.uv_handle_t?
  39. local function close_handle(handle)
  40. if handle and not handle:is_closing() then
  41. handle:close()
  42. end
  43. end
  44. --- @class vim.SystemObj
  45. --- @field cmd string[]
  46. --- @field pid integer
  47. --- @field private _state vim.SystemState
  48. --- @field wait fun(self: vim.SystemObj, timeout?: integer): vim.SystemCompleted
  49. --- @field kill fun(self: vim.SystemObj, signal: integer|string)
  50. --- @field write fun(self: vim.SystemObj, data?: string|string[])
  51. --- @field is_closing fun(self: vim.SystemObj): boolean
  52. local SystemObj = {}
  53. --- @param state vim.SystemState
  54. --- @return vim.SystemObj
  55. local function new_systemobj(state)
  56. return setmetatable({
  57. cmd = state.cmd,
  58. pid = state.pid,
  59. _state = state,
  60. }, { __index = SystemObj })
  61. end
  62. --- @param signal integer|string
  63. function SystemObj:kill(signal)
  64. self._state.handle:kill(signal)
  65. end
  66. --- @package
  67. --- @param signal? vim.SystemSig
  68. function SystemObj:_timeout(signal)
  69. self._state.done = 'timeout'
  70. self:kill(signal or SIG.TERM)
  71. end
  72. local MAX_TIMEOUT = 2 ^ 31
  73. --- @param timeout? integer
  74. --- @return vim.SystemCompleted
  75. function SystemObj:wait(timeout)
  76. local state = self._state
  77. local done = vim.wait(timeout or state.timeout or MAX_TIMEOUT, function()
  78. return state.result ~= nil
  79. end, nil, true)
  80. if not done then
  81. -- Send sigkill since this cannot be caught
  82. self:_timeout(SIG.KILL)
  83. vim.wait(timeout or state.timeout or MAX_TIMEOUT, function()
  84. return state.result ~= nil
  85. end, nil, true)
  86. end
  87. return state.result
  88. end
  89. --- @param data string[]|string|nil
  90. function SystemObj:write(data)
  91. local stdin = self._state.stdin
  92. if not stdin then
  93. error('stdin has not been opened on this object')
  94. end
  95. if type(data) == 'table' then
  96. for _, v in ipairs(data) do
  97. stdin:write(v)
  98. stdin:write('\n')
  99. end
  100. elseif type(data) == 'string' then
  101. stdin:write(data)
  102. elseif data == nil then
  103. -- Shutdown the write side of the duplex stream and then close the pipe.
  104. -- Note shutdown will wait for all the pending write requests to complete
  105. -- TODO(lewis6991): apparently shutdown doesn't behave this way.
  106. -- (https://github.com/neovim/neovim/pull/17620#discussion_r820775616)
  107. stdin:write('', function()
  108. stdin:shutdown(function()
  109. close_handle(stdin)
  110. end)
  111. end)
  112. end
  113. end
  114. --- @return boolean
  115. function SystemObj:is_closing()
  116. local handle = self._state.handle
  117. return handle == nil or handle:is_closing() or false
  118. end
  119. --- @param output? uv.read_start.callback|false
  120. --- @param text? boolean
  121. --- @return uv.uv_stream_t? pipe
  122. --- @return uv.read_start.callback? handler
  123. --- @return string[]? data
  124. local function setup_output(output, text)
  125. if output == false then
  126. return
  127. end
  128. local bucket --- @type string[]?
  129. local handler --- @type uv.read_start.callback
  130. if type(output) == 'function' then
  131. handler = output
  132. else
  133. bucket = {}
  134. handler = function(err, data)
  135. if err then
  136. error(err)
  137. end
  138. if text and data then
  139. bucket[#bucket + 1] = data:gsub('\r\n', '\n')
  140. else
  141. bucket[#bucket + 1] = data
  142. end
  143. end
  144. end
  145. local pipe = assert(uv.new_pipe(false))
  146. --- @type uv.read_start.callback
  147. local function handler_with_close(err, data)
  148. handler(err, data)
  149. if data == nil then
  150. pipe:read_stop()
  151. pipe:close()
  152. end
  153. end
  154. return pipe, handler_with_close, bucket
  155. end
  156. --- @param input? string|string[]|boolean
  157. --- @return uv.uv_stream_t?
  158. --- @return string|string[]?
  159. local function setup_input(input)
  160. if not input then
  161. return
  162. end
  163. local towrite --- @type string|string[]?
  164. if type(input) == 'string' or type(input) == 'table' then
  165. towrite = input
  166. end
  167. return assert(uv.new_pipe(false)), towrite
  168. end
  169. --- @return table<string,string>
  170. local function base_env()
  171. local env = vim.fn.environ() --- @type table<string,string>
  172. env['NVIM'] = vim.v.servername
  173. env['NVIM_LISTEN_ADDRESS'] = nil
  174. return env
  175. end
  176. --- uv.spawn will completely overwrite the environment
  177. --- when we just want to modify the existing one, so
  178. --- make sure to prepopulate it with the current env.
  179. --- @param env? table<string,string|number>
  180. --- @param clear_env? boolean
  181. --- @return string[]?
  182. local function setup_env(env, clear_env)
  183. if clear_env then
  184. return env
  185. end
  186. --- @type table<string,string|number>
  187. env = vim.tbl_extend('force', base_env(), env or {})
  188. local renv = {} --- @type string[]
  189. for k, v in pairs(env) do
  190. renv[#renv + 1] = string.format('%s=%s', k, tostring(v))
  191. end
  192. return renv
  193. end
  194. local is_win = vim.fn.has('win32') == 1
  195. local M = {}
  196. --- @param cmd string
  197. --- @param opts uv.spawn.options
  198. --- @param on_exit fun(code: integer, signal: integer)
  199. --- @param on_error fun()
  200. --- @return uv.uv_process_t, integer
  201. local function spawn(cmd, opts, on_exit, on_error)
  202. if is_win then
  203. local cmd1 = vim.fn.exepath(cmd)
  204. if cmd1 ~= '' then
  205. cmd = cmd1
  206. end
  207. end
  208. local handle, pid_or_err = uv.spawn(cmd, opts, on_exit)
  209. if not handle then
  210. on_error()
  211. error(pid_or_err)
  212. end
  213. return handle, pid_or_err --[[@as integer]]
  214. end
  215. --- @param timeout integer
  216. --- @param cb fun()
  217. --- @return uv.uv_timer_t
  218. local function timer_oneshot(timeout, cb)
  219. local timer = assert(uv.new_timer())
  220. timer:start(timeout, 0, function()
  221. timer:stop()
  222. timer:close()
  223. cb()
  224. end)
  225. return timer
  226. end
  227. --- @param state vim.SystemState
  228. --- @param code integer
  229. --- @param signal integer
  230. --- @param on_exit fun(result: vim.SystemCompleted)?
  231. local function _on_exit(state, code, signal, on_exit)
  232. close_handle(state.handle)
  233. close_handle(state.stdin)
  234. close_handle(state.timer)
  235. -- #30846: Do not close stdout/stderr here, as they may still have data to
  236. -- read. They will be closed in uv.read_start on EOF.
  237. local check = assert(uv.new_check())
  238. check:start(function()
  239. for _, pipe in pairs({ state.stdin, state.stdout, state.stderr }) do
  240. if not pipe:is_closing() then
  241. return
  242. end
  243. end
  244. check:stop()
  245. check:close()
  246. if state.done == nil then
  247. state.done = true
  248. end
  249. if (code == 0 or code == 1) and state.done == 'timeout' then
  250. -- Unix: code == 0
  251. -- Windows: code == 1
  252. code = 124
  253. end
  254. local stdout_data = state.stdout_data
  255. local stderr_data = state.stderr_data
  256. state.result = {
  257. code = code,
  258. signal = signal,
  259. stdout = stdout_data and table.concat(stdout_data) or nil,
  260. stderr = stderr_data and table.concat(stderr_data) or nil,
  261. }
  262. if on_exit then
  263. on_exit(state.result)
  264. end
  265. end)
  266. end
  267. --- @param state vim.SystemState
  268. local function _on_error(state)
  269. close_handle(state.handle)
  270. close_handle(state.stdin)
  271. close_handle(state.stdout)
  272. close_handle(state.stderr)
  273. close_handle(state.timer)
  274. end
  275. --- Run a system command
  276. ---
  277. --- @param cmd string[]
  278. --- @param opts? vim.SystemOpts
  279. --- @param on_exit? fun(out: vim.SystemCompleted)
  280. --- @return vim.SystemObj
  281. function M.run(cmd, opts, on_exit)
  282. vim.validate('cmd', cmd, 'table')
  283. vim.validate('opts', opts, 'table', true)
  284. vim.validate('on_exit', on_exit, 'function', true)
  285. opts = opts or {}
  286. local stdout, stdout_handler, stdout_data = setup_output(opts.stdout, opts.text)
  287. local stderr, stderr_handler, stderr_data = setup_output(opts.stderr, opts.text)
  288. local stdin, towrite = setup_input(opts.stdin)
  289. --- @type vim.SystemState
  290. local state = {
  291. done = false,
  292. cmd = cmd,
  293. timeout = opts.timeout,
  294. stdin = stdin,
  295. stdout = stdout,
  296. stdout_data = stdout_data,
  297. stderr = stderr,
  298. stderr_data = stderr_data,
  299. }
  300. --- @diagnostic disable-next-line:missing-fields
  301. state.handle, state.pid = spawn(cmd[1], {
  302. args = vim.list_slice(cmd, 2),
  303. stdio = { stdin, stdout, stderr },
  304. cwd = opts.cwd,
  305. --- @diagnostic disable-next-line:assign-type-mismatch
  306. env = setup_env(opts.env, opts.clear_env),
  307. detached = opts.detach,
  308. hide = true,
  309. }, function(code, signal)
  310. _on_exit(state, code, signal, on_exit)
  311. end, function()
  312. _on_error(state)
  313. end)
  314. if stdout and stdout_handler then
  315. stdout:read_start(stdout_handler)
  316. end
  317. if stderr and stderr_handler then
  318. stderr:read_start(stderr_handler)
  319. end
  320. local obj = new_systemobj(state)
  321. if towrite then
  322. obj:write(towrite)
  323. obj:write(nil) -- close the stream
  324. end
  325. if opts.timeout then
  326. state.timer = timer_oneshot(opts.timeout, function()
  327. if state.handle and state.handle:is_active() then
  328. obj:_timeout()
  329. end
  330. end)
  331. end
  332. return obj
  333. end
  334. return M