CSV.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | CSV reading and writing for dataframes. A strict, single-pass
RFC 4180 scanner (cassava-compatible) parses fields into typed column
builders; ragged rows are padded with nulls and extra fields dropped.
-}
module DataFrame.IO.CSV (
    -- * Reading
    readCsv,
    readTsv,
    readCsvWithOpts,
    readSeparated,
    readCsvWithSchema,
    CsvReader,
    decodeSeparated,
    fromCsv,
    fromCsvBytes,

    -- * Options
    module DataFrame.IO.CSV.Internal.Options,

    -- * Writing
    writeCsv,
    writeTsv,
    writeSeparated,

    -- * Helpers
    stripQuotes,
) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Text.IO as TIO

import Control.Exception (SomeException, catch)
import Data.Maybe (fromMaybe)
import DataFrame.IO.CSV.Internal.Options
import DataFrame.IO.CSV.Internal.Read (decodeCsvStrict)
import DataFrame.Internal.DataFrame (DataFrame (..), toSeparated)
import DataFrame.Schema (Schema, elements)

{- | Read CSV file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readCsv ".\/data\/taxi.csv"

@
-}
readCsv :: FilePath -> IO DataFrame
readCsv = readSeparated defaultReadOptions

type CsvReader = Schema -> FilePath -> IO DataFrame

{- | Schema-driven CSV reader. Coerces each column to the type declared
in 'Schema'; columns absent from the schema fall back to inference.

@
import qualified DataFrame as D
df <- D.readCsvWithSchema schema "input.csv"
@
-}
readCsvWithSchema :: CsvReader
readCsvWithSchema schema =
    readSeparated
        defaultReadOptions
            { typeSpec =
                SpecifyTypes
                    (M.toList (elements schema))
                    (typeSpec defaultReadOptions)
            }

{- | Read CSV file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readCsvWithOpts ".\/data\/taxi.csv" (D.defaultReadOptions { dateFormat = "%d/%-m/%-Y" })

@
-}
readCsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame
readCsvWithOpts = readSeparated

{- | Read TSV (tab separated) file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readTsv ".\/data\/taxi.tsv"

@
-}
readTsv :: FilePath -> IO DataFrame
readTsv = readSeparated (defaultReadOptions{columnSeparator = '\t'})

{- | Read text file with specified delimiter into a dataframe.

==== __Example__
@
ghci> D.readSeparated (D.defaultReadOptions { columnSeparator = ';' }) ".\/data\/taxi.txt"

@
-}
readSeparated :: ReadOptions -> FilePath -> IO DataFrame
readSeparated opts path = do
    let stripUtf8Bom b = fromMaybe b (BS.stripPrefix "\xEF\xBB\xBF" b)
    csvData <- stripUtf8Bom <$> BS.readFile path
    decodeCsvStrict opts csvData

{- | Decode in-memory CSV bytes into a dataframe. The result is fully
forced. (Note: unlike 'readSeparated', no UTF-8 BOM is stripped.)
-}
decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame
decodeSeparated opts csvData = decodeCsvStrict opts (BL.toStrict csvData)

writeCsv :: FilePath -> DataFrame -> IO ()
writeCsv = writeSeparated ','

writeTsv :: FilePath -> DataFrame -> IO ()
writeTsv = writeSeparated '\t'

writeSeparated ::
    -- | Separator
    Char ->
    -- | Path to write to
    FilePath ->
    DataFrame ->
    IO ()
writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df)

-- | Parse a CSV string into a DataFrame using default options.
fromCsv :: String -> IO (Either String DataFrame)
fromCsv s = do
    let bs = BL.fromStrict (TE.encodeUtf8 (T.pack s))
    (Right <$> decodeSeparated defaultReadOptions bs)
        `catch` (\(e :: SomeException) -> pure (Left (show e)))

-- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options.
fromCsvBytes :: BL.ByteString -> IO DataFrame
fromCsvBytes = decodeSeparated defaultReadOptions

stripQuotes :: T.Text -> T.Text
stripQuotes txt =
    case T.uncons txt of
        Just ('"', rest) ->
            case T.unsnoc rest of
                Just (middle, '"') -> middle
                _ -> txt
        _ -> txt