_system.lua 9.6 KB

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