feat: add type annotations for plenary.async

Author: delphinus

lua/plenary/async/api.lua

@@ -1,5 +1,8 @@
 local util = require "plenary.async.util"
 
+---@alias PlenaryAsyncApi table<string, PlenaryAsyncFunction>
+
+---@type PlenaryAsyncApi
 return setmetatable({}, {
   __index = function(t, k)
     return function(...)

lua/plenary/async/async.lua

@@ -4,13 +4,20 @@ local errors = require "plenary.errors"
 local traceback_error = errors.traceback_error
 local f = require "plenary.functional"
 
+---@class PlenaryAsyncAsync
 local M = {}
 
+---@param fn any
+---@return boolean
 local function is_callable(fn)
   return type(fn) == "function" or (type(fn) == "table" and type(getmetatable(fn)["__call"]) == "function")
 end
 
 ---because we can't store varargs
+---@param step function
+---@param thread thread
+---@param callback? function
+---@param ... any
 local function callback_or_next(step, thread, callback, ...)
   local stat = f.first(...)
 
@@ -33,12 +40,12 @@ local function callback_or_next(step, thread, callback, ...)
 end
 
 ---Executes a future with a callback when it is done
----@param async_function Future: the future to execute
----@param callback function: the callback to call when done
-local execute = function(async_function, callback, ...)
-  assert(is_callable(async_function), "type error :: expected func")
+---@param func function the future to execute
+---@param callback? function the callback to call when done
+local execute = function(func, callback, ...)
+  assert(is_callable(func), "type error :: expected func")
 
-  local thread = co.create(async_function)
+  local thread = co.create(func)
 
   local step
   step = function(...)
@@ -48,31 +55,41 @@ local execute = function(async_function, callback, ...)
   step(...)
 end
 
+---A function including async logic
+---@alias PlenaryAsyncFunction async fun(...): ...
+
 local add_leaf_function
 do
   ---A table to store all leaf async functions
+  ---@type table<PlenaryAsyncFunction, integer>
   _PlenaryLeafTable = setmetatable({}, {
     __mode = "k",
   })
 
+  ---@param async_func PlenaryAsyncFunction
+  ---@param argc integer
   add_leaf_function = function(async_func, argc)
     assert(_PlenaryLeafTable[async_func] == nil, "Async function should not already be in the table")
     _PlenaryLeafTable[async_func] = argc
   end
 
+  ---@param async_func PlenaryAsyncFunction
+  ---@return boolean
   function M.is_leaf_function(async_func)
     return _PlenaryLeafTable[async_func] ~= nil
   end
 
+  ---@param async_func PlenaryAsyncFunction
+  ---@return integer
   function M.get_leaf_function_argc(async_func)
     return _PlenaryLeafTable[async_func]
   end
 end
 
 ---Creates an async function with a callback style function.
----@param func function: A callback style function to be converted. The last argument must be the callback.
----@param argc number: The number of arguments of func. Must be included.
----@return function: Returns an async function
+---@param func function A callback style function to be converted. The last argument must be the callback.
+---@param argc integer The number of arguments of func. Must be included.
+---@return PlenaryAsyncFunction leaf Returns an leaf
 M.wrap = function(func, argc)
   if not is_callable(func) then
     traceback_error("type error :: expected func, got " .. type(func))
@@ -82,6 +99,7 @@ M.wrap = function(func, argc)
     traceback_error("type error :: expected number, got " .. type(argc))
   end
 
+  ---@type PlenaryAsyncFunction
   local function leaf(...)
     local nargs = select("#", ...)
 
@@ -99,13 +117,13 @@ end
 
 ---Use this to either run a future concurrently and then do something else
 ---or use it to run a future with a callback in a non async context
----@param async_function function
+---@param func PlenaryAsyncFunction|function
 ---@param callback function
-M.run = function(async_function, callback)
-  if M.is_leaf_function(async_function) then
-    async_function(callback)
+M.run = function(func, callback)
+  if M.is_leaf_function(func) then
+    func(callback)
   else
-    execute(async_function, callback)
+    execute(func, callback)
   end
 end
 

lua/plenary/async/control.lua

@@ -2,22 +2,26 @@ local a = require "plenary.async.async"
 local Deque = require("plenary.async.structs").Deque
 local tbl = require "plenary.tbl"
 
+---@class PlenaryAsyncControl
 local M = {}
 
+---@class PlenaryAsyncCondvar
+---@field handles (fun(): nil)[]
 local Condvar = {}
 Condvar.__index = Condvar
 
----@class Condvar
----@return Condvar
+---@return PlenaryAsyncCondvar
 function Condvar.new()
   return setmetatable({ handles = {} }, Condvar)
 end
 
 ---`blocks` the thread until a notification is received
+---@param self PlenaryAsyncCondvar
+---@param callback fun(): nil
 Condvar.wait = a.wrap(function(self, callback)
   -- not calling the callback will block the coroutine
   table.insert(self.handles, callback)
-end, 2)
+end, 2) --[[@as async fun(): nil]]
 
 ---notify everyone that is waiting on this Condvar
 function Condvar:notify_all()
@@ -52,12 +56,14 @@ end
 
 M.Condvar = Condvar
 
+---@class PlenaryAsyncSemaphore
+---@field handles (fun(permit: PlenaryAsyncPermit): nil)[]
+---@field permits integer
 local Semaphore = {}
 Semaphore.__index = Semaphore
 
----@class Semaphore
----@param initial_permits number: the number of permits that it can give out
----@return Semaphore
+---@param initial_permits integer the number of permits that it can give out
+---@return PlenaryAsyncSemaphore
 function Semaphore.new(initial_permits)
   vim.validate {
     initial_permits = {
@@ -79,6 +85,8 @@ end
 ---permit:forget()
 ---when a permit can be acquired returns it
 ---call permit:forget() to forget the permit
+---@param self PlenaryAsyncSemaphore
+---@param callback fun(permit: PlenaryAsyncPermit): nil
 Semaphore.acquire = a.wrap(function(self, callback)
   if self.permits > 0 then
     self.permits = self.permits - 1
@@ -87,8 +95,10 @@ Semaphore.acquire = a.wrap(function(self, callback)
     return
   end
 
+  ---@class PlenaryAsyncPermit
   local permit = {}
 
+  ---@param self_permit PlenaryAsyncPermit
   permit.forget = function(self_permit)
     self.permits = self.permits + 1
 
@@ -99,16 +109,17 @@ Semaphore.acquire = a.wrap(function(self, callback)
   end
 
   callback(permit)
-end, 2)
+end, 2) --[[@as async fun(self: PlenaryAsyncSemaphore): PlenaryAsyncPermit]]
 
 M.Semaphore = Semaphore
 
+---@class PlenaryAsyncControlChannel
 M.channel = {}
 
 ---Creates a oneshot channel
 ---returns a sender and receiver function
 ---the sender is not async while the receiver is
----@return function, function
+---@return fun(...): nil tx, async fun(): ... rx
 M.channel.oneshot = function()
   local val = nil
   local saved_callback = nil
@@ -147,20 +158,26 @@ M.channel.oneshot = function()
       if is_single then
         return callback(val)
       else
-        return callback(tbl.unpack(val))
+        return callback(tbl.unpack(val --[[@as table]]))
       end
     else
       saved_callback = callback
     end
-  end, 1)
+  end, 1) --[[@as async fun(): ...]]
 
   return sender, receiver
 end
 
+---@class PlenaryAsyncCounterTx
+---@field send fun(): nil
+
+---@class PlenaryAsyncCounterRx
+---@field recv async fun(): nil
+---@field last async fun(): nil
+
 ---A counter channel.
 ---Basically a channel that you want to use only to notify and not to send any actual values.
----@return function: sender
----@return function: receiver
+---@return PlenaryAsyncCounterTx tx, PlenaryAsyncCounterRx rx
 M.channel.counter = function()
   local counter = 0
   local condvar = Condvar.new()
@@ -191,9 +208,15 @@ M.channel.counter = function()
   return Sender, Receiver
 end
 
+---@class PlenaryAsyncMpscTx
+---@field send fun(...: any): nil
+
+---@class PlenaryAsyncMpscRx
+---@field recv async fun(): ...
+---@field last async fun(): ...
+
 ---A multiple producer single consumer channel
----@return table
----@return table
+---@return PlenaryAsyncMpscTx, PlenaryAsyncMpscRx
 M.channel.mpsc = function()
   local deque = Deque.new()
   local condvar = Condvar.new()

lua/plenary/async/init.lua

@@ -12,6 +12,13 @@ local lookups = {
   control = "plenary.async.control",
 }
 
+---@class PlenaryAsync: PlenaryAsyncAsync
+---@field api PlenaryAsyncApi
+---@field control PlenaryAsyncControl
+---@field lsp PlenaryAsyncLsp
+---@field tests PlenaryAsyncTests
+---@field util PlenaryAsyncUtil
+---@field uv PlenaryAsyncUv
 local exports = setmetatable(require "plenary.async.async", {
   __index = function(t, k)
     local require_path = lookups[k]

lua/plenary/async/lsp.lua

@@ -1,12 +1,29 @@
 local a = require "plenary.async.async"
 
+---@class PlenaryAsyncLsp
 local M = {}
 
+---@alias ClientRequestIds table<integer, integer>
+
 ---This will be deprecated because the callback can be called multiple times.
 ---This will give a coroutine error because the coroutine will be resumed multiple times.
 ---Please use buf_request_all instead.
+---@type async fun(bufnr: integer, method: string, params?: table, handler?: lsp.Handler): ClientRequestIds, function
 M.buf_request = a.wrap(vim.lsp.buf_request, 4)
 
+---@alias BufRequestAllHandler fun(results: table<integer, { error: lsp.ResponseError, result: any }>)
+
+---Sends an async request for all active clients attached to the buffer and executes the `handler`
+---callback with the combined result.
+---
+---* param bufnr (integer) Buffer handle, or 0 for current.
+---* param method (string) LSP method name
+---* param params (table|nil) Parameters to send to the server
+---* param handler fun(results: table<integer, {error: lsp.ResponseError, result: any}>) (function)
+--- Handler called after all requests are completed. Server results are passed as
+--- a `client_id:result` map.
+---* return function cancel Function that cancels all requests.
+---@type async fun(bufnr: integer, method: string, params?: table, handler: BufRequestAllHandler): function
 M.buf_request_all = a.wrap(vim.lsp.buf_request_all, 4)
 
 return M

lua/plenary/async/structs.lua

@@ -1,12 +1,15 @@
+---@class PlenaryAsyncStructs
 local M = {}
 
+---A double ended queue
+---@class PlenaryDeque
+---@field first integer
+---@field last integer
+---@field [integer] any
 local Deque = {}
 Deque.__index = Deque
 
----@class Deque
----A double ended queue
----
----@return Deque
+---@return PlenaryDeque
 function Deque.new()
   -- the indexes are created with an offset so that the indices are consequtive
   -- otherwise, when both pushleft and pushright are used, the indices will have a 1 length hole in the middle
@@ -62,13 +65,13 @@ function Deque:is_empty()
 end
 
 ---returns the number of elements of the deque
----@return number
+---@return integer
 function Deque:len()
   return self.last - self.first + 1
 end
 
 ---returns and iterator of the indices and values starting from the left
----@return function
+---@return fun(): integer?, any?
 function Deque:ipairs_left()
   local i = self.first
 
@@ -85,7 +88,7 @@ function Deque:ipairs_left()
 end
 
 ---returns and iterator of the indices and values starting from the right
----@return function
+---@return fun(): integer?, any?
 function Deque:ipairs_right()
   local i = self.last
 

lua/plenary/async/tests.lua

@@ -1,23 +1,31 @@
 local util = require "plenary.async.util"
 
+---@class PlenaryAsyncTests
 local M = {}
 
+---@param s string
+---@param async_func PlenaryAsyncFunction
 M.describe = function(s, async_func)
   describe(s, async_func)
 end
 
+---@param s string
+---@param async_func PlenaryAsyncFunction
 M.it = function(s, async_func)
   it(s, util.will_block(async_func, tonumber(vim.env.PLENARY_TEST_TIMEOUT)))
 end
 
+---@param async_func PlenaryAsyncFunction
 M.pending = function(async_func)
   pending(async_func)
 end
 
+---@param async_func PlenaryAsyncFunction
 M.before_each = function(async_func)
   before_each(util.will_block(async_func))
 end
 
+---@param async_func PlenaryAsyncFunction
 M.after_each = function(async_func)
   after_each(util.will_block(async_func))
 end