ftcsv.lua

local ftcsv = {
    _VERSION = 'ftcsv 1.1.0',
    _DESCRIPTION = 'CSV library for Lua',
    _URL         = 'https://github.com/FourierTransformer/ftcsv',
    _LICENSE     = [[
        The MIT License (MIT)

        Copyright (c) 2016 Shakil Thakur

        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:

        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.

        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
    ]]
}

-- lua 5.1 load compat
local M = {}
if type(jit) == 'table' or _ENV then
    M.load = _G.load
else
    M.load = loadstring
end

-- perf
local sbyte = string.byte
local ssub = string.sub

-- luajit specific speedups
-- luajit performs faster with iterating over string.byte,
-- whereas vanilla lua performs faster with string.find
if type(jit) == 'table' then
    -- finds the end of an escape sequence
    function M.findClosingQuote(i, inputLength, inputString, quote, doubleQuoteEscape)
        -- local doubleQuoteEscape = doubleQuoteEscape
        local currentChar, nextChar = sbyte(inputString, i), nil
        while i <= inputLength do
            -- print(i)
            nextChar = sbyte(inputString, i+1)

            -- this one deals with " double quotes that are escaped "" within single quotes "
            -- these should be turned into a single quote at the end of the field
            if currentChar == quote and nextChar == quote then
                doubleQuoteEscape = true
                i = i + 2
                currentChar = sbyte(inputString, i)

            -- identifies the escape toggle
            elseif currentChar == quote and nextChar ~= quote then
                -- print("exiting", i-1)
                return i-1, doubleQuoteEscape
            else
                i = i + 1
                currentChar = nextChar
            end
        end
    end

else
    -- vanilla lua closing quote finder
    function M.findClosingQuote(i, inputLength, inputString, quote, doubleQuoteEscape)
        local firstCharIndex = 1
        local firstChar, iChar = nil, nil
        repeat
            firstCharIndex, i = inputString:find('".?', i+1)
            firstChar = sbyte(inputString, firstCharIndex)
            iChar = sbyte(inputString, i)
            -- nextChar = string.byte(inputString, i+1)
            -- print("HI", offset, i)
            -- print(firstChar, iChar)
            if firstChar == quote and iChar == quote then
                doubleQuoteEscape = true
            end
        until iChar ~= quote
        if i == nil then
            return inputLength-1, doubleQuoteEscape
        end
        -- print("exiting", i-2)
        return i-2, doubleQuoteEscape
    end

end

-- load an entire file into memory
local function loadFile(textFile)
    local file = io.open(textFile, "r")
    if not file then error("File not found at " .. textFile) end
    local allLines = file:read("*all")
    file:close()
    return allLines
end

-- creates a new field and adds it to the main table
local function createField(inputString, quote, fieldStart, i, doubleQuoteEscape)
    local field
    -- so, if we just recently de-escaped, we don't want the trailing \"
    if sbyte(inputString, i-1) == quote then
        -- print("Skipping last \"")
        field = ssub(inputString, fieldStart, i-2)
    else
        field = ssub(inputString, fieldStart, i-1)
    end
    if doubleQuoteEscape then
        -- print("QUOTE REPLACE")
        -- print(line[fieldNum])
        field = field:gsub('""', '"')
    end
    return field
end

-- main function used to parse
local function parseString(inputString, inputLength, delimiter, i, headerField, fieldsToKeep)

    -- keep track of my chars!
    local currentChar, nextChar = sbyte(inputString, i), nil
    local skipChar = 0
    local field
    local fieldStart = i
    local fieldNum = 1
    local lineNum = 1
    local doubleQuoteEscape = false
    local exit = false

    --bytes
    local CR = sbyte("\r")
    local LF = sbyte("\n")
    local quote = sbyte('"')
    local delimiterByte = sbyte(delimiter)

    local assignValue
    local outResults
    -- the headers haven't been set yet.
    -- aka this is the first run!
    if headerField == nil then
        -- print("this is for headers")
        headerField = {}
        assignValue = function()
            headerField[fieldNum] = field
            return true
        end
    else
        -- print("this is for magic")
        outResults = {}
        outResults[1] = {}
        assignValue = function()
            outResults[lineNum][headerField[fieldNum]] = field
        end
    end

    while i <= inputLength do
        -- go by two chars at a time! currentChar is set at the bottom.
        -- currentChar = string.byte(inputString, i)
        nextChar = sbyte(inputString, i+1)
        -- print(i, string.char(currentChar), string.char(nextChar))

        -- empty string
        if currentChar == quote and nextChar == quote then
            -- print("EMPTY STRING")
            skipChar = 1
            fieldStart = i + 2
            -- print("fs+2:", fieldStart)

        -- identifies the escape toggle
        elseif currentChar == quote and nextChar ~= quote then
            -- print("ESCAPE TOGGLE")
            fieldStart = i + 1
            i, doubleQuoteEscape = M.findClosingQuote(i+1, inputLength, inputString, quote, doubleQuoteEscape)
            -- print("I VALUE", i, doubleQuoteEscape)
            skipChar = 1

        -- create some fields if we can!
        elseif currentChar == delimiterByte then
            -- create the new field
            -- print(headerField[fieldNum])
            if fieldsToKeep == nil or fieldsToKeep[headerField[fieldNum]] then
                field = createField(inputString, quote, fieldStart, i, doubleQuoteEscape)
            -- print("FIELD", field, "FIELDEND", headerField[fieldNum], lineNum)
            -- outResults[headerField[fieldNum]][lineNum] = field
                assignValue()
            end
            doubleQuoteEscape = false

            fieldNum = fieldNum + 1
            fieldStart = i + 1
            -- print("fs+1:", fieldStart)
        -- end

        -- newline?!
        elseif ((currentChar == CR and nextChar == LF) or currentChar == LF) then
            if fieldsToKeep == nil or fieldsToKeep[headerField[fieldNum]] then
                -- create the new field
                field = createField(inputString, quote, fieldStart, i, doubleQuoteEscape)

                -- outResults[headerField[fieldNum]][lineNum] = field
                exit = assignValue()
                if exit then
                    if (currentChar == CR and nextChar == LF) then
                        return headerField, i + 1
                    else
                        return headerField, i
                    end
                end
            end
            doubleQuoteEscape = false

            -- determine how line ends
            if (currentChar == CR and nextChar == LF) then
                -- print("CRLF DETECTED")
                skipChar = 1
                fieldStart = fieldStart + 1
                -- print("fs:", fieldStart)
            end

            -- incrememnt for new line
            lineNum = lineNum + 1
            outResults[lineNum] = {}
            fieldNum = 1
            fieldStart = i + 1 + skipChar
            -- print("fs:", fieldStart)

        end

        i = i + 1 + skipChar
        if (skipChar > 0) then
            currentChar = sbyte(inputString, i)
        else
            currentChar = nextChar
        end
        skipChar = 0
    end

    -- create last new field
    if fieldsToKeep == nil or fieldsToKeep[headerField[fieldNum]] then
        field = createField(inputString, quote, fieldStart, i, doubleQuoteEscape)
        assignValue()
    end

    -- clean up last line if it's weird (this happens when there is a CRLF newline at end of file)
    -- doing a count gets it to pick up the oddballs
    local finalLineCount = 0
    for _, _ in pairs(outResults[lineNum]) do
        finalLineCount = finalLineCount + 1
    end
    local initialLineCount = 0
    for _, _ in pairs(outResults[1]) do
        initialLineCount = initialLineCount + 1
    end
    -- print("Final/Initial", finalLineCount, initialLineCount)
    if finalLineCount ~= initialLineCount then
        outResults[lineNum] = nil
    end

    return outResults
end

-- runs the show!
function ftcsv.parse(inputFile, delimiter, options)
    -- delimiter MUST be one character
    assert(#delimiter == 1 and type(delimiter) == "string", "the delimiter must be of string type and exactly one character")

    -- OPTIONS yo
    local header = true
    local rename
    local fieldsToKeep = nil
    local loadFromString = false
    local headerFunc
    if options then
        if options.headers ~= nil then
            assert(type(options.headers) == "boolean", "ftcsv only takes the boolean 'true' or 'false' for the optional parameter 'headers' (default 'true'). You passed in '" .. tostring(options.headers) .. "' of type '" .. type(options.headers) .. "'.")
            header = options.headers
        end
        if options.rename ~= nil then
            assert(type(options.rename) == "table", "ftcsv only takes in a key-value table for the optional parameter 'rename'. You passed in '" .. tostring(options.rename) .. "' of type '" .. type(options.rename) .. "'.")
            rename = options.rename
        end
        if options.fieldsToKeep ~= nil then
            assert(type(options.fieldsToKeep) == "table", "ftcsv only takes in a list (as a table) for the optional parameter 'fieldsToKeep'. You passed in '" .. tostring(options.fieldsToKeep) .. "' of type '" .. type(options.fieldsToKeep) .. "'.")
            local ofieldsToKeep = options.fieldsToKeep
            if ofieldsToKeep ~= nil then
                fieldsToKeep = {}
                for j = 1, #ofieldsToKeep do
                    fieldsToKeep[ofieldsToKeep[j]] = true
                end
            end
            if header == false then
                assert(next(rename) ~= nil, "ftcsv can only have fieldsToKeep for header-less files when they have been renamed. Please add the 'rename' option and try again.")
            end
        end
        if options.loadFromString ~= nil then
            assert(type(options.loadFromString) == "boolean", "ftcsv only takes a boolean value for optional parameter 'loadFromString'. You passed in '" .. tostring(options.loadFromString) .. "' of type '" .. type(options.loadFromString) .. "'.")
            loadFromString = options.loadFromString
        end
        if options.headerFunc ~= nil then
            assert(type(options.headerFunc) == "function", "ftcsv only takes a function value for optional parameter 'headerFunc'. You passed in '" .. tostring(options.headerFunc) .. "' of type '" .. type(options.headerFunc) .. "'.")
            headerFunc = options.headerFunc
        end
    end

    -- handle input via string or file!
    local inputString
    if loadFromString then
        inputString = inputFile
    else
        inputString = loadFile(inputFile)
    end
    local inputLength = #inputString

    -- parse through the headers!
    local headerField, i = parseString(inputString, inputLength, delimiter, 0)
    i = i + 1 -- start at the next char

    -- for files where there aren't headers!
    if header == false then
        i = 0
        for j = 1, #headerField do
            headerField[j] = j
        end
    end

    -- rename fields as needed!
    if rename then
        -- basic rename (["a" = "apple"])
        for j = 1, #headerField do
            if rename[headerField[j]] then
                -- print("RENAMING", headerField[j], rename[headerField[j]])
                headerField[j] = rename[headerField[j]]
            end
        end
        -- files without headers, but with a rename need to be handled too!
        if #rename > 0 then
            for j = 1, #rename do
                headerField[j] = rename[j]
            end
        end
    end

    -- apply some sweet header manuipulation
    if headerFunc then
        for j = 1, #headerField do
            headerField[j] = headerFunc(headerField[j])
        end
    end

    local output = parseString(inputString, inputLength, delimiter, i, headerField, fieldsToKeep)
    return output
end

-- a function that delimits " to "", used by the writer
local function delimitField(field)
    field = tostring(field)
    if field:find('"') then
        return field:gsub('"', '""')
    else
        return field
    end
end

-- a function that compiles some lua code to quickly print out the csv
local function writer(inputTable, dilimeter, headers)
    -- they get re-created here if they need to be escaped so lua understands it based on how
    -- they came in
    for i = 1, #headers do
        if inputTable[1][headers[i]] == nil then
            error("the field '" .. headers[i] .. "' doesn't exist in the table")
        end
        if headers[i]:find('"') then
            headers[i] = headers[i]:gsub('"', '\\"')
        end
    end

    local outputFunc = [[
        local state, i = ...
        local d = state.delimitField
        i = i + 1;
        if i > state.tableSize then return nil end;
        return i, '"' .. d(state.t[i]["]] .. table.concat(headers, [["]) .. '"]] .. dilimeter .. [["' .. d(state.t[i]["]]) .. [["]) .. '"\r\n']]

    -- print(outputFunc)

    local state = {}
    state.t = inputTable
    state.tableSize = #inputTable
    state.delimitField = delimitField

    return M.load(outputFunc), state, 0

end

-- takes the values from the headers in the first row of the input table
local function extractHeaders(inputTable)
    local headers = {}
    for key, _ in pairs(inputTable[1]) do
        headers[#headers+1] = key
    end

    -- lets make the headers alphabetical
    table.sort(headers)

    return headers
end

-- turns a lua table into a csv
-- works really quickly with luajit-2.1, because table.concat life
function ftcsv.encode(inputTable, delimiter, options)
    local output = {}

    -- dilimeter MUST be one character
    assert(#delimiter == 1 and type(delimiter) == "string", "the delimiter must be of string type and exactly one character")

    -- grab the headers from the options if they are there
    local headers = nil
    if options then
        if options.fieldsToKeep ~= nil then
            assert(type(options.fieldsToKeep) == "table", "ftcsv only takes in a list (as a table) for the optional parameter 'fieldsToKeep'. You passed in '" .. tostring(options.headers) .. "' of type '" .. type(options.headers) .. "'.")
            headers = options.fieldsToKeep
        end
    end
    if headers == nil then
        headers = extractHeaders(inputTable)
    end

    -- newHeaders are needed if there are quotes within the header
    -- because they need to be escaped
    local newHeaders = {}
    for i = 1, #headers do
        if headers[i]:find('"') then
            newHeaders[i] = headers[i]:gsub('"', '""')
        else
            newHeaders[i] = headers[i]
        end
    end
    output[1] = '"' .. table.concat(newHeaders, '","') .. '"\r\n'

    -- add each line by line.
    for i, line in writer(inputTable, delimiter, headers) do
        output[i+1] = line
    end
    return table.concat(output)
end

return ftcsv