Add writeFile for Builder (#408)

* Add writeFile for `Builder` #406 * Clarify imports in documentation Co-authored-by: Bodigrim <andrew.lelechenko@gmail.com> * Add test for `writeFile` * Evaluate read bytestring before closing file Co-authored-by: Bodigrim <andrew.lelechenko@gmail.com>
haskell · Aug 5, 2021 · 1010c9e · 1010c9e
1 parent 8e8ee17
commit 1010c9e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 1 deletion.
diff --git a/Data/ByteString/Builder.hs b/Data/ByteString/Builder.hs
@@ -190,6 +190,7 @@ module Data.ByteString.Builder
       -- about fine-tuning them.
     , toLazyByteString
     , hPutBuilder
+    , writeFile
 
       -- * Creating Builders
 
@@ -254,13 +255,15 @@ module Data.ByteString.Builder
 
     ) where
 
+import           Prelude hiding (writeFile)
+
 import           Data.ByteString.Builder.Internal
 import qualified Data.ByteString.Builder.Prim  as P
 import qualified Data.ByteString.Lazy.Internal as L
 import           Data.ByteString.Builder.ASCII
 
 import           Data.String (IsString(..))
-import           System.IO (Handle)
+import           System.IO (Handle, IOMode(..), withBinaryFile)
 import           Foreign
 import           GHC.Base (unpackCString#, unpackCStringUtf8#,
                            unpackFoldrCString#, build)
@@ -293,6 +296,17 @@ toLazyByteString = toLazyByteStringWith
 hPutBuilder :: Handle -> Builder -> IO ()
 hPutBuilder h = hPut h . putBuilder
 
+modifyFile :: IOMode -> FilePath -> Builder -> IO ()
+modifyFile mode f bld = withBinaryFile f mode (`hPutBuilder` bld)
+
+-- | Write a 'Builder' to a file.
+--
+-- Similarly to 'hPutBuilder', this function is more efficient than 
+-- using 'Data.ByteString.Lazy.hPut' . 'toLazyByteString' with a file handle. 
+--
+-- @since 0.11.2.0
+writeFile :: FilePath -> Builder -> IO ()
+writeFile = modifyFile WriteMode
 
 ------------------------------------------------------------------------------
 -- Binary encodings

diff --git a/tests/builder/Data/ByteString/Builder/Tests.hs b/tests/builder/Data/ByteString/Builder/Tests.hs
@@ -14,6 +14,8 @@
 
 module Data.ByteString.Builder.Tests (tests) where
 
+import           Prelude hiding (writeFile)
+
 import           Control.Applicative
 import           Control.Monad (unless, void)
 import           Control.Monad.Trans.State (StateT, evalStateT, evalState, put, get)
@@ -59,6 +61,7 @@ tests =
   , testHandlePutBuilderChar8
   , testPut
   , testRunBuilder
+  , testWriteFile
   ] ++
   testsEncodingToBuilder ++
   testsBinary ++
@@ -165,6 +168,32 @@ testHandlePutBuilderChar8 =
         unless success (error msg)
         return success
 
+testWriteFile :: TestTree
+testWriteFile =
+    testProperty "writeFile" testRecipe
+  where
+    testRecipe :: Recipe -> Property
+    testRecipe recipe =
+        ioProperty $ do
+            (tempFile, tempH) <- openTempFile "." "test-builder-writeFile.tmp"
+            hClose tempH
+            let b = fst $ recipeComponents recipe
+            writeFile tempFile b
+            lbs <- L.readFile tempFile
+            _ <- evaluate (L.length $ lbs)
+            removeFile tempFile
+            let lbsRef = toLazyByteString b
+            -- report
+            let msg =
+                    unlines
+                        [ "recipe:   " ++ show recipe
+                        , "via file: " ++ show lbs
+                        , "direct :  " ++ show lbsRef
+                        ]
+                success = lbs == lbsRef
+            unless success (error msg)
+            return success
+
 removeFile :: String -> IO ()
 removeFile fn = void $ withCString fn c_unlink