In Lua, is there a way to document a function type including a description for its parameters?

Question

I'm doing some maintenance in my NeoVim config, almost entirely written in Lua.
For convenience, I created a mini-lib of useful tools to ease this configuration.
I have several friends who keep a close eye on my NeoVim config, or even using it, and for that reason I like my toolings to be well documented to ease their experience with it.
Coming mainly from a TypeScript background, you can imagine that I like to go deep with it.

In this example, I have a module that reduces the hassle of setting up mappings, also making it more readable.

--- Set a keymap using Vim's API
---@param desc  string     The descritpion that will be shown for your mapping when calling :map
---@param mode  MapModeStr The mode that your keymap will work in
---@param lhs   string     A NeoVim keymap string
---@param rhs   MapRHS     The action yout keymap will trigger
---@param opts? MapOptions A table of options
local set = function(desc, mode, lhs, rhs, opts)
  opts = opts or {}
  local finalopts = vim.tbl_deep_extend("keep", { desc = desc }, opts)
  vim.keymap.set(mode, lhs, rhs, finalopts)
end

return {
  set = set,

  --- Set a keymap for Insert mode
  ---@param desc  string     The descritpion that will be shown for your mapping when calling :map
  ---@param lhs   string     A NeoVim keymap string
  ---@param rhs   MapRHS     The action yout keymap will trigger
  ---@param opts? MapOptions A table of options
  i = function(desc, lhs, rhs, opts) set(desc, "i", lhs, rhs, opts) end,

  --- Set a keymap for Normal mode
  ---@param desc  string     The descritpion that will be shown for your mapping when calling :map
  ---@param lhs   string     A NeoVim keymap string
  ---@param rhs   MapRHS     The action yout keymap will trigger
  ---@param opts? MapOptions A table of options
  n = function(desc, lhs, rhs, opts) set(desc, "n", lhs, rhs, opts) end,

  ---@param desc  string     The descritpion that will be shown for your mapping when calling :map
  ---@param lhs   string     A NeoVim keymap string
  ---@param rhs   MapRHS     The action yout keymap will trigger
  ---@param opts? MapOptions A table of options
  o = function(desc, lhs, rhs, opts) end,

  --- Set a keymap for Visual and Select mode
  ---@param desc  string     The descritpion that will be shown for your mapping when calling :map
  ---@param lhs   string     A NeoVim keymap string
  ---@param rhs   MapRHS     The action yout keymap will trigger
  ---@param opts? MapOptions A table of options
  x = function(desc, lhs, rhs, opts) set(desc, "x", lhs, rhs, opts) end,

  --- Set a keymap for Visual mode
  ---@param desc  string     The descritpion that will be shown for your mapping when calling :map
  ---@param lhs   string     A NeoVim keymap string
  ---@param rhs   MapRHS     The action yout keymap will trigger
  ---@param opts? MapOptions A table of options
  v = function(desc, lhs, rhs, opts) set(desc, "v", lhs, rhs, opts) end,
}

To reduce copypasta, I would like to "typedef" the mapper functions that are exposed to my users in a meta file with the other types I've defined for this module. I know this is something I can do:

---@meta

...

---@alias MapperFunction fun(desc: string, lhs: string, rhs: MapRHS, opts: MapOptions?)

And then I could simply ---@type MapperFunction on them, but the issue then is that I lose the descriptions of what each parameter does, and I would like to avoid that.

The best I could come up with was some spicy little curry (functional programming rocks!). It works, but I wanted to know if there was a way to do it with only typings.

An annoying side effect is that I lose some of the completion annotations.

So, is there a way I could simply slap a type on top of my functions and retain all IntelliSense, or am I doomed to use one of those two techniques ?

Thanks!

Answer 1

I am not sure, if the below suggestion would be useful to you in your particular environment, but it may be interesting to Lua programmers in general.

It is a 'decorator' that allows to create Lua functions with typed arguments and returns. Function specification is a string, e.g.:

[[
Squares its argument.
@param number x Number to square
@return number The squared number
]]

Nullable types (starting with ?) are allowed; so are multiple types, e.g. number|string.

local unpack = unpack or table.unpack

local tags = {
    author = 'An author of the module or file',
    copyright = 'The copyright notice of the module or file. LuaDoc adds a © sign between the label (Copyright) and the given text (e.g. 2004-2007 Kepler Project)',
    field = 'Describe a table field definition',
    param = 'Describe function parameters. It requires the name of the parameter and its description',
    release = 'Free format string to describe the module or file release',
    ['return'] = 'Describe a returning value of the function. Since Lua can return multiple values, this tag should appear more than once',
    see = 'Refers to other descriptions of functions or tables',
    usage = 'Describe the usage of the function or variable',
    class = 'If LuaDoc cannot infer the type of documentation (function, table or module definition), the programmer can specify it explicitly',
    description = 'The description of the function or table. This is usually infered automatically',
    name = 'The name of the function or table definition. This is usually infered from the code analysis, and the programmer does not need to define it. If LuaDoc can infer the name of the function automatically it\'s even not recomended to define the name explicitly, to avoid redundancy',
    type = 'Set "type" of a table.',
    precondition = 'A function that should return true for passed parameters',
    postcondition = 'A function that should return true for returned parameters'
}

local tag_fields = { param = 2 }

-- Makes an array of possible types out of string '[?]type1|[?]type2|...|[?]typen'.
local function type_spec (spec)
    local variants = {}
    for type in spec:gmatch '[^|]+' do
        if type:sub (1, 1) == '?' then
            variants ['nil'] = true
            type = type:sub (2)
        end
        variants [type] = true
    end
    return variants
end

local function get_tag (line)
    local tag, def = line:match '^%s*-*%s*@(%S+)(.*)$'
    if def and tags [tag] then
        local pattern = '^' .. ('%s+(%S+)'):rep (tag_fields[tag] or 1) .. '(.*)$'
        local args = { def:match (pattern) }
        if #args > 0 then
            return tag, args
        end
    end
    return 'comment', line
end

-- Returns an or-separated list of types, with nullable types prefixed by '?'.
local function serialise_types (type_set)
    local types = {}
    for type, _ in pairs (type_set) do
        types [#types + 1] = type
    end
    return table.concat (types, ' or ')
end

-- Convert @param / @return tag to precondition / postcondition:
local function tag2condition (tag, spec, no)
    local expected, name, desc = type_spec (spec [1]), #spec > 1 and spec [2] or nil, spec [#spec]
    local serialised = serialise_types (expected)
    local message
        = 'Type mismatch. ' .. (tag == 'return' and 'Returned value' or 'Parameter') .. ' '
        .. tostring (no) .. (name and ' (' .. name .. ': ' .. (desc or '(no description)') .. ')' or '')
        .. ': ' .. serialised .. ' is expected, '
        .. 'but %s %s is ' .. (tag == 'return' and 'returned' or 'supplied')
    return function (...)
        local value = ({...}) [no]
        local actual = type (value) == 'table' and value.type or type (value)
        if not expected [actual] then
            error (message:format (actual, tostring (value)))
        end
        return true
    end
end

-- Parse the docs:
local function parse_doc (doc)
    local parsed = {}
    if type (doc) == 'table' then
        parsed = doc
    else
        for line in doc:gmatch '[^\n]+' do
            local tag, args = get_tag (line)
            parsed [tag] = parsed [tag] or {}
            parsed [tag] [#parsed [tag] + 1] = args
        end
    end
    -- Convert @param to preconditions and @return to postconditions:
    for tag, class in pairs { param = 'precondition', ['return'] = 'postcondition' } do
        for no, spec in ipairs (parsed [tag]) do
            parsed [class] = parsed [class] or {}
            parsed [class] [#parsed [class] + 1] = tag2condition (tag, spec, no)
        end
    end
    return parsed
end

-- Checks an array of values against an array of expected type specifications.
-- Aborts with error, if any of values is of wrong type.
local function check_conditions (conditions, ...)
    for _, condition in ipairs (conditions) do
        condition (...)
    end
    return true
end

-- Wraps documented with a callable table with type checks for parameters and returns as defined by spec.
local function document (doc, documented)
    local parsed = parse_doc (doc)
    parsed.doc, parsed.raw = doc, documented
    -- Use f.validate_params (...) to get true/false, [error text].
    parsed.validate_params = function (...)
        return pcall (check_conditions, parsed.precondition or {}, ...)
    end
    local t, mt = {}, {}
    if type (documented) == 'function' then
        t = parsed
        mt.__call = function (self, ...) -- simply call f(...).
            check_conditions (parsed.precondition or {}, ...) -- check parameters against the specification.
            local returned = { documented (...) } -- actually call the function.
            check_conditions (parsed.postcondition or {}, unpack (returned))
            return unpack (returned) -- return the validated return values.
        end
    elseif type (documented) == 'table' then
        t = documented
        t.type = t.type or (parsed.type [1] or {}) [1] or 'table'
        for key, values in pairs (parsed) do
            t [key] = t [key] or values
        end
    end
    return setmetatable (t, mt)
end

-- That's how functions are decorated to make them typed.
local squared = document ([[
Squares its argument
--- @param number x Number to square
--- @return number The squared number
]],
    function (x) return x ^ 2 end
)

print ('Description: ', squared.doc)

for _, arg in ipairs { 2, 'two' } do
    print ('x = ', arg, 'validation: ', squared.validate_params (arg))
    print ('x = ', arg, 'squared (x) = ', squared (arg))
end