Decorators.hs

{-# LANGUAGE TemplateHaskellQuotes #-}

module Decorators (myDecorator) where

import Wraps (Decorator, wraps)

-- | Example from https://docs.python.org/3/library/functools.html#functools.wraps
myDecorator :: Decorator
myDecorator = wraps $ \f -> [e|
    do
        putStrLn "Calling decorated function"
        $(f) |]

Main.hs

{-# LANGUAGE TemplateHaskell #-}

module Main (main) where

import Decorators (myDecorator)

myDecorator [d|
    example1 :: IO ()
    example1 = putStrLn "Called example1 function" |]

myDecorator [d|
    example2 :: Int -> IO ()
    example2 n = putStrLn $ "Called example2 function with " ++ show n |]

main :: IO ()
main = do
    example1
    example2 42

Wraps.hs

module Wraps (Decorator, wraps) where

import Control.Lens
import Language.Haskell.TH
import Language.Haskell.TH.Lens
 
type Decorator = DecsQ -> DecsQ

wraps :: (ExpQ -> ExpQ) -> Decorator
wraps fExpQ = (>>= mapM wrapDec) where
    wrapDec :: Dec -> DecQ
    wrapDec = mapMOf (decBodies . bodyExps) (fExpQ . return)

    decBodies = _FunD . _2 . each . clauseBody `failing` _ValD . _2
    bodyExps = _GuardedB . each . _2 `failing` _NormalB