{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}

{- |
Module      : DataFrame.TH.Records
License     : MIT

Record-based 'DataFrame' splices — the IO-agnostic core of the
'DataFrame.TH' family. Splices that read CSV / Parquet files at compile
time live in @DataFrame.TH.CSV@ (in @dataframe-csv-th@) and
@DataFrame.TH.Parquet@ (in @dataframe-parquet-th@).
-}
module DataFrame.TH.Records (
    -- * Declare one binding per column
    declareColumns,
    declareColumnsWithPrefix,
    declareColumnsWithPrefix',

    -- * Type-string parser (exposed for testing)
    typeFromString,
) where

import Data.Function (on)
import Data.Functor ((<&>))
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Text as T

import Control.Monad (forM)
import Language.Haskell.TH
import qualified Language.Haskell.TH.Syntax as TH

import DataFrame.Functions (sanitize)
import DataFrame.Internal.Column (columnTypeString)
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    unsafeGetColumn,
 )
import DataFrame.Internal.Expression (Expr)
import DataFrame.Operators (col)
import Prelude as P

typeFromString :: [String] -> Q Type
typeFromString :: [String] -> Q Type
typeFromString [] = String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"No type specified"
typeFromString [String
t0] = do
    let t :: String
t = String -> String
trim String
t0
    case String -> Maybe String
stripBrackets String
t of
        Just String
inner -> [String] -> Q Type
typeFromString [String
inner] Q Type -> (Type -> Type) -> Q Type
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Type -> Type -> Type
AppT Type
ListT
        Maybe String
Nothing
            | String
t String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"Text" Bool -> Bool -> Bool
|| String
t String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"Data.Text.Text" Bool -> Bool -> Bool
|| String
t String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"T.Text" ->
                Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Type
ConT ''T.Text)
            | Bool
otherwise -> do
                Maybe Name
m <- String -> Q (Maybe Name)
lookupTypeName String
t
                case Maybe Name
m of
                    Just Name
tyName -> Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Type
ConT Name
tyName)
                    Maybe Name
Nothing -> String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$ String
"Unsupported type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
t0
typeFromString [String
tycon, String
t1] = Type -> Type -> Type
AppT (Type -> Type -> Type) -> Q Type -> Q (Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String] -> Q Type
typeFromString [String
tycon] Q (Type -> Type) -> Q Type -> Q Type
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [String] -> Q Type
typeFromString [String
t1]
typeFromString [String
tycon, String
t1, String
t2] =
    (\Type
outer Type
a Type
b -> Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
outer Type
a) Type
b)
        (Type -> Type -> Type -> Type)
-> Q Type -> Q (Type -> Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String] -> Q Type
typeFromString [String
tycon]
        Q (Type -> Type -> Type) -> Q Type -> Q (Type -> Type)
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [String] -> Q Type
typeFromString [String
t1]
        Q (Type -> Type) -> Q Type -> Q Type
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [String] -> Q Type
typeFromString [String
t2]
typeFromString [String]
s = String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$ String
"Unsupported types: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords [String]
s

trim :: String -> String
trim :: String -> String
trim = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
' ') (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
forall a. [a] -> [a]
reverse (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
' ') (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
forall a. [a] -> [a]
reverse

stripBrackets :: String -> Maybe String
stripBrackets :: String -> Maybe String
stripBrackets String
s =
    case String
s of
        (Char
'[' : String
rest)
            | Bool -> Bool
P.not (String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
rest) Bool -> Bool -> Bool
&& String -> Char
forall a. HasCallStack => [a] -> a
last String
rest Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
']' ->
                String -> Maybe String
forall a. a -> Maybe a
Just (String -> String
forall a. HasCallStack => [a] -> [a]
init String
rest)
        String
_ -> Maybe String
forall a. Maybe a
Nothing

{- | Splice a binding for every column of @df@, named after the column.
Column names that are not valid Haskell identifiers are sanitized
(see 'DataFrame.Functions.sanitize').
-}
declareColumns :: DataFrame -> DecsQ
declareColumns :: DataFrame -> DecsQ
declareColumns = Maybe Text -> DataFrame -> DecsQ
declareColumnsWithPrefix' Maybe Text
forall a. Maybe a
Nothing

-- | Like 'declareColumns' but prefixes every binding name with @prefix_@.
declareColumnsWithPrefix :: T.Text -> DataFrame -> DecsQ
declareColumnsWithPrefix :: Text -> DataFrame -> DecsQ
declareColumnsWithPrefix Text
prefix = Maybe Text -> DataFrame -> DecsQ
declareColumnsWithPrefix' (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
prefix)

-- | Like 'declareColumnsWithPrefix' but takes an optional prefix.
declareColumnsWithPrefix' :: Maybe T.Text -> DataFrame -> DecsQ
declareColumnsWithPrefix' :: Maybe Text -> DataFrame -> DecsQ
declareColumnsWithPrefix' Maybe Text
prefix DataFrame
df =
    let
        names :: [Text]
names = (((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst ([(Text, Int)] -> [Text])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Text, Int)] -> [(Text, Int)])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)])
-> (DataFrame -> Map Text Int) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Map Text Int
columnIndices) DataFrame
df
        types :: [String]
types = (Text -> String) -> [Text] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> String
columnTypeString (Column -> String) -> (Text -> Column) -> Text -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> DataFrame -> Column
`unsafeGetColumn` DataFrame
df)) [Text]
names
        specs :: [(Text, Text, String)]
specs =
            (Text -> String -> (Text, Text, String))
-> [Text] -> [String] -> [(Text, Text, String)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
                ( \Text
colName String
type_ ->
                    ( Text
colName
                    , Text -> (Text -> Text) -> Maybe Text -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" (Text -> Text
sanitize (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_")) Maybe Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
sanitize Text
colName
                    , String
type_
                    )
                )
                [Text]
names
                [String]
types
     in
        ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [(Text, Text, String)]
-> ((Text, Text, String) -> DecsQ) -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Text, Text, String)]
specs (((Text, Text, String) -> DecsQ) -> Q [[Dec]])
-> ((Text, Text, String) -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \(Text
raw, Text
nm, String
tyStr) -> do
            Type
ty <- [String] -> Q Type
typeFromString (String -> [String]
words String
tyStr)
            let n :: Name
n = String -> Name
mkName (Text -> String
T.unpack Text
nm)
            Dec
sig <- Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
n [t|Expr $(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty)|]
            Dec
val <- Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|col $(Text -> Q Exp
forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
forall (m :: * -> *). Quote m => Text -> m Exp
TH.lift Text
raw)|]) []
            [Dec] -> DecsQ
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
sig, Dec
val]