Merge pull request #85 from szeiger/streaming-file-output

Write output directly to disk with -o

sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala

@@ -1,6 +1,7 @@
 package sjsonnet
 
-import java.io.{BufferedOutputStream, InputStream, OutputStreamWriter, PrintStream, StringWriter}
+import java.io.{BufferedOutputStream, InputStream, OutputStreamWriter, PrintStream, StringWriter, Writer}
+import java.nio.charset.StandardCharsets
 import java.nio.file.NoSuchFileException
 
 import sjsonnet.Cli.Config
@@ -75,10 +76,10 @@ object SjsonnetMain {
 
     result match{
       case Left(err) =>
-        if (!err.isEmpty) System.err.println(err)
+        if (!err.isEmpty) stderr.println(err)
         1
       case Right(str) =>
-        if (!str.isEmpty) System.out.println(str)
+        if (!str.isEmpty) stdout.println(str)
         0
     }
   }
@@ -105,15 +106,25 @@ object SjsonnetMain {
       config.preserveOrder
     )
 
-    def writeFile(f: os.RelPath, contents: String): Either[String, Unit] = {
-      Try(os.write.over(os.Path(f, wd), contents, createFolders = config.createDirs))
-        .toEither
-        .left
-        .map{
-          case e: NoSuchFileException => s"open $f: no such file or directory"
-          case e => e.toString
-        }
-    }
+    def handleWriteFile[T](f: => T): Either[String, T] =
+      Try(f).toEither.left.map{
+        case e: NoSuchFileException => s"open $f: no such file or directory"
+        case e => e.toString
+      }
+
+    def writeFile(f: os.RelPath, contents: String): Either[String, Unit] =
+      handleWriteFile(os.write.over(os.Path(f, wd), contents, createFolders = config.createDirs))
+
+    def writeToFile[U](f: os.RelPath, contents: Writer => Either[String, U]): Either[String, Unit] =
+      handleWriteFile(os.write.over.outputStream(os.Path(f, wd), createFolders = config.createDirs)).flatMap { out =>
+        try {
+          val buf = new BufferedOutputStream(out)
+          val wr = new OutputStreamWriter(buf, StandardCharsets.UTF_8)
+          val u = contents(wr)
+          wr.flush()
+          u.map(_ => ())
+        } finally out.close()
+      }
 
     (config.multi, config.yamlStream) match {
       case (Some(multiPath), _) =>
@@ -167,22 +178,20 @@ object SjsonnetMain {
               "whose elements hold the JSON for each document in the stream.")
         }
       case _ =>
-        val materialized = interp.interpret0(
+        def materialize(wr: Writer) = interp.interpret0(
           os.read(path),
           OsPath(path),
-          new Renderer(indent = config.indent)
+          new Renderer(wr, indent = config.indent)
         )
         config.outputFile match{
-          case None => materialized.map(_.toString)
+          case None =>
+            materialize(new StringWriter).map(_.toString)
           case Some(f) =>
             val filePath = os.FilePath(f) match{
               case _: os.Path => os.Path(f).relativeTo(os.pwd)
               case _ => os.RelPath(f)
             }
-            for{
-              materializedStr <- materialized
-              _ <- writeFile(filePath, materializedStr.toString)
-            } yield ""
+            writeToFile(filePath, materialize(_)).map(_ => "")
         }
     }
   }

sjsonnet/src/sjsonnet/Renderer.scala

@@ -1,5 +1,5 @@
 package sjsonnet
-import java.io.StringWriter
+import java.io.{StringWriter, Writer}
 import java.util.regex.Pattern
 
 import upickle.core.{ArrVisitor, ObjVisitor}
@@ -13,7 +13,7 @@ import ujson.BaseRenderer
   * - Custom printing of empty dictionaries and arrays
   *
   */
-class Renderer(out: StringWriter = new java.io.StringWriter(),
+class Renderer(out: Writer = new java.io.StringWriter(),
                indent: Int = -1) extends BaseRenderer(out, indent){
   var newlineBuffered = false
   override def visitFloat64(d: Double, index: Int) = {
@@ -40,15 +40,15 @@ class Renderer(out: StringWriter = new java.io.StringWriter(),
     newlineBuffered = false
     commaBuffered = false
   }
-  override def visitArray(length: Int, index: Int) = new ArrVisitor[StringWriter, StringWriter] {
+  override def visitArray(length: Int, index: Int) = new ArrVisitor[Writer, Writer] {
     var empty = true
     flushBuffer()
     out.append('[')
     newlineBuffered = true
 
     depth += 1
     def subVisitor = Renderer.this
-    def visitValue(v: StringWriter, index: Int): Unit = {
+    def visitValue(v: Writer, index: Int): Unit = {
       empty = false
       flushBuffer()
       commaBuffered = true
@@ -65,7 +65,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(),
     }
   }
 
-  override def visitObject(length: Int, index: Int) = new ObjVisitor[StringWriter, StringWriter] {
+  override def visitObject(length: Int, index: Int) = new ObjVisitor[Writer, Writer] {
     var empty = true
     flushBuffer()
     out.append('{')
@@ -78,7 +78,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(),
       flushBuffer()
       out.append(colonSnippet)
     }
-    def visitValue(v: StringWriter, index: Int): Unit = {
+    def visitValue(v: Writer, index: Int): Unit = {
       commaBuffered = true
     }
     def visitEnd(index: Int) = {
@@ -212,7 +212,7 @@ class YamlRenderer(out: StringWriter = new java.io.StringWriter(), indentArrayIn
   }
 }
 
-class PythonRenderer(out: StringWriter = new java.io.StringWriter(),
+class PythonRenderer(out: Writer = new java.io.StringWriter(),
                      indent: Int = -1) extends BaseRenderer(out, indent){
 
   override def visitNull(index: Int) = {

sjsonnet/test/src-jvm/sjsonnet/MainTests.scala

@@ -0,0 +1,32 @@
+package sjsonnet
+
+import java.io.{ByteArrayOutputStream, File, PrintStream}
+import java.util.Arrays
+
+import utest._
+
+object MainTests extends TestSuite {
+
+  val testSuiteRoot = os.pwd / "sjsonnet" / "test" / "resources" / "test_suite"
+
+  val tests = Tests {
+    // Compare writing to stdout with writing to a file
+    test("writeToFile") {
+      val source = (testSuiteRoot / "local.jsonnet").toString()
+      val outF = File.createTempFile("sjsonnet", ".json")
+      val out = new ByteArrayOutputStream()
+      val pout = new PrintStream(out)
+      SjsonnetMain.main0(Array(source), collection.mutable.Map.empty, System.in, pout, System.err, os.pwd, None)
+      pout.flush()
+      SjsonnetMain.main0(Array("-o", outF.getAbsolutePath, source), collection.mutable.Map.empty, System.in, System.out, System.err, os.pwd, None)
+      val stdoutBytes = out.toByteArray
+      val fileBytes = os.read(os.Path(outF)).getBytes
+      // stdout mode uses println so it has an extra platform-specific line separator at the end
+      val eol = System.getProperty("line.separator").getBytes
+
+      //println(stdoutBytes.map(_.toInt).mkString(","))
+      //println(fileBytes.map(_.toInt).mkString(","))
+      assert(Arrays.equals(fileBytes ++ eol, stdoutBytes))
+    }
+  }
+}