aesontrans.hs

{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
module AesonTrans (
  AesonTrans, fromField, maybeFromField, preTransJSON, postTransJSON,
  transFieldModifier)
where
import Data.Aeson
import Data.Aeson.Types
import Data.Text (Text)
import qualified Data.Text as Text
import Control.Monad
import Control.Monad.Trans.Maybe
import GHC.Records
import GHC.TypeLits
import Data.Proxy
import qualified Data.HashMap.Strict as HashMap

data AesonTrans a =
  SubField Text [Text] Text |
  TransAfter (AesonTrans a) (AesonTrans a) |
  SubFieldMaybe Text [Text] Text |
  NoTrans

instance Semigroup (AesonTrans a) where
  (<>) = TransAfter

instance Monoid (AesonTrans a) where
  mempty = NoTrans

fromField :: (HasField s a b, KnownSymbol s)
          => Proxy (s :: Symbol) -> [Text] -> Text -> AesonTrans a
fromField proxy = SubField (Text.pack $ symbolVal proxy)

maybeFromField :: (HasField s a (Maybe b), KnownSymbol s)
               => Proxy (s :: Symbol) -> [Text] -> Text -> AesonTrans a
maybeFromField proxy = SubFieldMaybe (Text.pack $ symbolVal proxy)

getRemove :: Object -> Text -> Maybe (Object, Value)
getRemove obj t =
  do val <- HashMap.lookup t obj
     pure (HashMap.delete t obj, val)
     
getRemovePath :: Object -> [Text] -> Text -> Parser (Maybe (Object, Value))
getRemovePath obj [] field = pure $ getRemove obj field
getRemovePath obj (path1:pathRest) field = runMaybeT $ do
  subObject <- MaybeT $ obj .:? path1
  (removed, val) <- MaybeT $ getRemovePath subObject pathRest field
  if HashMap.null removed
    then pure (HashMap.delete path1 obj, val)
    else pure (HashMap.insert path1 (Object removed) obj, val)

transParser :: String -> AesonTrans a -> Value -> Parser Value
transParser name (TransAfter t1 t2) =
  transParser name t1 >=> transParser name t2
transParser name (SubField to path from) =
  withObject name $ \obj -> do
  res <- getRemovePath obj path from
  case res of
    Nothing -> fail $ "field not found: " ++ Text.unpack from ++ " in path " ++
               show path
    Just (o, val) -> pure $ Object $ HashMap.insert to val o
transParser name (SubFieldMaybe to path from) =
  withObject name $ \obj -> do
  res <- getRemovePath obj path from
  case res of
    Nothing -> pure $ Object obj
    Just (o, val) -> pure $ Object $ HashMap.insert to val o
transParser _ NoTrans = pure

-- | transform a json value before parsing it.
preTransJSON :: String -> AesonTrans a -> (Value -> Parser a)
             -> (Value -> Parser a)
preTransJSON name trans parser = transParser name trans >=> parser

putPath :: Text -> Value -> [Text] -> Object -> Object
putPath to val [] obj = HashMap.insert to val obj
putPath to val (path:paths) obj = 
  let newMap = case HashMap.lookup path obj of
        Nothing -> HashMap.empty
        Just (Object mp) -> mp
        _ -> error "putPath: expected Object"
  in HashMap.insert to (Object $ putPath path val paths newMap) obj

toJSONTrans :: AesonTrans a -> Value -> Value
toJSONTrans trans val@(Object obj) =
  case trans of
    TransAfter t1 t2 -> toJSONTrans t2 . toJSONTrans t1 $ val
    SubField to path from -> case HashMap.lookup to obj of
      Nothing -> error "toJSONTrans: object not found"
      Just fieldVal ->
        Object $ putPath from fieldVal path $ HashMap.delete to obj
    SubFieldMaybe to path from -> case HashMap.lookup to obj of
      Nothing -> val
      Just fieldVal ->
        Object $ putPath from fieldVal path $ HashMap.delete to obj
    NoTrans -> val
toJSONTrans _ _ = error "toJSONTrans: expected an object"

-- | transform a json value after generating it
postTransJSON :: AesonTrans a -> (a -> Value) -> (a -> Value)
postTransJSON trans coder = toJSONTrans trans . coder

onText :: (String -> String) -> (Text -> Text)
onText f = Text.pack . f . Text.unpack

-- | apply a field modifier to the transformation
transFieldModifier :: (String -> String) -> AesonTrans a -> AesonTrans a
transFieldModifier m (TransAfter t1 t2) =
  TransAfter (transFieldModifier m t1) (transFieldModifier m t2)
transFieldModifier m (SubField to path from) =
  SubField (onText m to) path from
transFieldModifier m (SubFieldMaybe to path from) =
  SubFieldMaybe (onText m to) path from
transFieldModifier _ NoTrans = NoTrans