path.hs

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Lens
import Control.Monad (void)
import Data.Aeson
import Data.Aeson.Lens
import Data.Aeson.QQ.Simple
import Data.Char
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Void
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer

data PathStep
  = Key Text
  | Nth Int
  deriving (Show)

type Path = [PathStep]

parseStep :: Parsec Void Text PathStep
parseStep =
  choice
    [ Nth <$> between (char '[') (char ']') decimal,
      Key . Text.pack <$> (char '.' *> manyTill anySingle (lookAhead (void (char '.')) <|> eof <|> lookAhead (void (char '['))))
    ]

parseSteps = many parseStep <* eof

path :: AsValue t => Text -> Traversal' t Value
path p =
  case runParser parseSteps "" p of
    Left e -> error $ errorBundlePretty e
    Right ps -> _Value . foldr (.) id (fmap toT ps)
  where
    toT (Key k) = key k
    toT (Nth i) = nth i

main = do
  print $ [aesonQQ|{"a":[1, {"b": 2}]}|] ^? path ".a[1].b" . _Number