@@ -215,7 +215,7 @@ module Streamly.Streams.StreamD
215215 , zipWith
216216 , zipWithM
217217
218− -- * Comparisions
218+ -- * Comparisons
219219 , eqBy
220220 , cmpBy
221221
@@ -225,9 +225,19 @@ module Streamly.Streams.StreamD
225225
226226 -- * Transformation comprehensions
227227 , the
228+
229+ -- * Exceptions
230+ , before
231+ , after
232+ , bracket
233+ , onException
234+ , finally
235+ , handle
228236 )
229237where
230238
239+import Control.Exception (Exception)
240+import Control.Monad.Catch (MonadCatch)
231241import Control.Monad.IO.Class (MonadIO(..))
232242import Control.Monad.Trans (MonadTrans(lift))
233243import Data.Word (Word32)
@@ -243,6 +253,8 @@ import Prelude
243253 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
244254 reverse)
245255
256+import qualified Control.Monad.Catch as MC
257+
246258import Streamly.Mem.Array.Types (Array(..))
247259import Streamly.Fold.Types (Fold(..))
248260import Streamly.SVar (MonadAsync, defState, adaptState)
@@ -1663,6 +1675,100 @@ mapM_ m = drain . mapM m
16631675fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
16641676fromStreamD = K.fromStream . toStreamK
16651677
1678+------------------------------------------------------------------------------
1679+-- Exceptions
1680+------------------------------------------------------------------------------
1681+
1682+-- | Run a side effect before the stream yields its first element.
1683+before :: Monad m => m b -> Stream m a -> Stream m a
1684+before action (Stream step state) = Stream step' Nothing
1685+
1686+ where
1687+
1688+ {-# INLINE_LATE step' #-}
1689+ step' _ Nothing = action >> return (Skip (Just state))
1690+
1691+ step' gst (Just st) = do
1692+ res <- step gst st
1693+ case res of
1694+ Yield x s -> return $ Yield x (Just s)
1695+ Skip s -> return $ Skip (Just s)
1696+ Stop -> return Stop
1697+
1698+-- | Run a side effect whenever the stream stops normally.
1699+after :: Monad m => m b -> Stream m a -> Stream m a
1700+after action (Stream step state) = Stream step' state
1701+
1702+ where
1703+
1704+ {-# INLINE_LATE step' #-}
1705+ step' gst st = do
1706+ res <- step gst st
1707+ case res of
1708+ Yield x s -> return $ Yield x s
1709+ Skip s -> return $ Skip s
1710+ Stop -> action >> return Stop
1711+
1712+-- | Run a side effect whenever the stream aborts due to an exception. The
1713+-- exception is not caught simply rethrown.
1714+onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1715+onException action (Stream step state) = Stream step' state
1716+
1717+ where
1718+
1719+ {-# INLINE_LATE step' #-}
1720+ step' gst st = do
1721+ res <- step gst st `MC.onException` action
1722+ case res of
1723+ Yield x s -> return $ Yield x s
1724+ Skip s -> return $ Skip s
1725+ Stop -> action >> return Stop
1726+
1727+-- XXX check if a monolithic implementation would be more performant.
1728+--
1729+-- | Run a side effect whenever the stream stops normally or aborts due to an
1730+-- exception.
1731+finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1732+finally action xs = after action $ onException action xs
1733+
1734+-- | Run the first action before the stream starts and remember its output,
1735+-- generate a stream using the output, run the second action using the
1736+-- remembered value as an argument whenever the stream ends normally or due to
1737+-- an exception.
1738+bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1739+bracket bef aft bet = Stream step' Nothing
1740+
1741+ where
1742+
1743+ {-# INLINE_LATE step' #-}
1744+ step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1745+
1746+ step' gst (Just (Stream step state, v)) = do
1747+ res <- step gst state `MC.onException` aft v
1748+ case res of
1749+ Yield x s -> return $ Yield x (Just (Stream step s, v))
1750+ Skip s -> return $ Skip (Just (Stream step s, v))
1751+ Stop -> aft v >> return Stop
1752+
1753+-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1754+-- and the specified exception handler is run with the exception as argument.
1755+handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1756+handle f (Stream step state) = Stream step' (Just state)
1757+
1758+ where
1759+
1760+ {-# INLINE_LATE step' #-}
1761+ step' gst (Just st) = do
1762+ res <- MC.try $ step gst st
1763+ case res of
1764+ Left e -> f e >>= \x -> return (Yield x Nothing)
1765+ Right r -> case r of
1766+ Yield x s -> return $ Yield x (Just s)
1767+ Skip s -> return $ Skip (Just s)
1768+ Stop -> return Stop
1769+
1770+ step' _ Nothing = return Stop
1771+
16661772------------------------------------------------------------------------------
16671773-- Transformation by Folding (Scans)
16681774------------------------------------------------------------------------------