dotfiles/.config/nvim/custom/maj-peg/lua/maj-peg/init.lua

local Job = require("plenary.job")
local Path = require("plenary.path")
local namespace_name = "maj-peg"
local status = {
    setup = false,
    enabled = true,
    has_mypy = false,
    has_mypy_baseline = false
}

local setup = function(mypy_flags, mypybaseline_flags)
    local namespace = vim.api.nvim_create_namespace(namespace_name)
    local group = vim.api.nvim_create_augroup(namespace_name, { clear = true })
    status.has_mypy = vim.fn.system({ "which", "mypy" }) ~= ""
    status.has_mypy_baseline = vim.fn.system({ "which", "mypy-baseline" }) ~= ""
    if mypy_flags == nil then
        mypy_flags = { "--follow-imports", "silent" }
    end
    if mypybaseline_flags == nil then
        mypybaseline_flags = { "filter" }
    end
    vim.api.nvim_create_autocmd({ "FileType" },
        {
            pattern = "python",
            group = group,
            callback = function(outer_args)
                vim.api.nvim_create_autocmd({ "BufWritePost", "FileType" }, {
                    buffer = outer_args.buf,
                    group = group,
                    callback = function(args)
                        vim.diagnostic.reset(namespace, args.buf)
                        if status.enabled == false then return end
                        local command = "mypy " ..
                            table.concat(mypy_flags, " ") .. " " ..
                            vim.api.nvim_buf_get_name(args.buf)
                        if status.has_mypy_baseline then
                            command = command ..
                                " | mypy-baseline " ..
                                table.concat(mypybaseline_flags, " ")
                        end
                        local path = Path:new(vim.api.nvim_buf_get_name(args.buf))
                            :make_relative()
                        local diagnostics = {}
                        Job:new({
                            command = "zsh",
                            args = { "-c", command },
                            on_stdout = function(_, line)
                                if line == nil or line:sub(1, path:len()) ~= path then
                                    return
                                end

                                line = line:sub(path:len() + 1)
                                local lnum = line:match(':%d+:')
                                line = line:gsub(':%d: ', '', 1)
                                if lnum == nil then return end
                                lnum = lnum:gsub(':', '')
                                lnum = tonumber(lnum) - 1

                                local severity = vim.diagnostic.severity.INFO
                                local severity_text = line:match('%a+: ')
                                if severity_text == nil or severity_text == "hint: " then
                                    severity = vim.diagnostic.severity.HINT
                                elseif severity_text == "error: " then
                                    severity = vim.diagnostic.severity.ERROR
                                elseif severity_text == "warn: " then
                                    severity = vim.diagnostic.severity.WARN
                                elseif severity_text == "info: " then
                                    severity = vim.diagnostic.severity.INFO
                                end
                                line = line:gsub('.+: ', '', 1)

                                table.insert(
                                    diagnostics,
                                    {
                                        lnum = tonumber(lnum),
                                        col = 0,
                                        source = namespace_name,
                                        severity = severity,
                                        message = line
                                    })
                            end,
                            on_exit = function()
                                vim.schedule(function()
                                    vim.diagnostic.set(namespace,
                                        args.buf, diagnostics)
                                end)
                            end,
                        }):start()
                    end
                })
            end
        })
    status.setup = true
end

return {
    setup = setup,
    toggle = function()
        status.enabled = not status.enabled;
    end,
    status = function()
        return status
    end
}