Implement quote_keys and c_document_end for yaml functions (#231)

Forked from PR #143

---------

Co-authored-by: Yi Jin <yi.jin@databricks.com>

Author: stephenamar-db

sjsonnet/src/sjsonnet/BaseCharRenderer.scala

@@ -4,16 +4,16 @@ package sjsonnet
 // with some private definitions made accessible to subclasses
 
 import ujson._
+
 import scala.annotation.switch
-import upickle.core.{ArrVisitor, ObjVisitor}
+import upickle.core.{ArrVisitor, ObjVisitor, Visitor}
 class BaseCharRenderer[T <: upickle.core.CharOps.Output]
 (out: T,
  indent: Int = -1,
  escapeUnicode: Boolean = false,
  newline: Array[Char] = Array('\n')) extends JsVisitor[T, T]{
   protected[this] val elemBuilder = new upickle.core.CharBuilder
-  protected[this] val unicodeCharBuilder = new upickle.core.CharBuilder()
-  def flushCharBuilder() = {
+  def flushCharBuilder(): Unit = {
     elemBuilder.writeOutToIfLongerThan(out, if (depth == 0) 0 else 1000)
   }
 
@@ -22,25 +22,29 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
 
   protected[this] var commaBuffered = false
 
-  def flushBuffer() = {
+  def flushBuffer(): Unit = {
     if (commaBuffered) {
       commaBuffered = false
       elemBuilder.append(',')
       renderIndent()
     }
   }
-  def visitArray(length: Int, index: Int) = new ArrVisitor[T, T] {
+
+  def visitArray(length: Int, index: Int): ArrVisitor[T, T] = new ArrVisitor[T, T] {
     flushBuffer()
     elemBuilder.append('[')
 
     depth += 1
     renderIndent()
-    def subVisitor = BaseCharRenderer.this
+
+    def subVisitor: Visitor[T, T] = BaseCharRenderer.this
+
     def visitValue(v: T, index: Int): Unit = {
       flushBuffer()
       commaBuffered = true
     }
-    def visitEnd(index: Int) = {
+
+    def visitEnd(index: Int): T = {
       commaBuffered = false
       depth -= 1
       renderIndent()
@@ -50,21 +54,26 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     }
   }
 
-  def visitObject(length: Int, index: Int) = new ObjVisitor[T, T] {
+  def visitObject(length: Int, index: Int): ObjVisitor[T, T] = new ObjVisitor[T, T] {
     flushBuffer()
     elemBuilder.append('{')
     depth += 1
     renderIndent()
-    def subVisitor = BaseCharRenderer.this
-    def visitKey(index: Int) = BaseCharRenderer.this
+
+    def subVisitor: Visitor[T, T] = BaseCharRenderer.this
+
+    def visitKey(index: Int): Visitor[T, T] = BaseCharRenderer.this
+
     def visitKeyValue(s: Any): Unit = {
       elemBuilder.append(':')
       if (indent != -1) elemBuilder.append(' ')
     }
+
     def visitValue(v: T, index: Int): Unit = {
       commaBuffered = true
     }
-    def visitEnd(index: Int) = {
+
+    def visitEnd(index: Int): T = {
       commaBuffered = false
       depth -= 1
       renderIndent()
@@ -74,7 +83,7 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     }
   }
 
-  def visitNull(index: Int) = {
+  def visitNull(index: Int): T = {
     flushBuffer()
     elemBuilder.ensureLength(4)
     elemBuilder.appendUnsafe('n')
@@ -85,7 +94,7 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     out
   }
 
-  def visitFalse(index: Int) = {
+  def visitFalse(index: Int): T = {
     flushBuffer()
     elemBuilder.ensureLength(5)
     elemBuilder.appendUnsafe('f')
@@ -97,7 +106,7 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     out
   }
 
-  def visitTrue(index: Int) = {
+  def visitTrue(index: Int): T = {
     flushBuffer()
     elemBuilder.ensureLength(4)
     elemBuilder.appendUnsafe('t')
@@ -108,7 +117,7 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     out
   }
 
-  def visitFloat64StringParts(s: CharSequence, decIndex: Int, expIndex: Int, index: Int) = {
+  def visitFloat64StringParts(s: CharSequence, decIndex: Int, expIndex: Int, index: Int): T = {
     flushBuffer()
     elemBuilder.ensureLength(s.length())
     var i = 0
@@ -121,7 +130,7 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
     out
   }
 
-  override def visitFloat64(d: Double, index: Int) = {
+  override def visitFloat64(d: Double, index: Int): T = {
     d match{
       case Double.PositiveInfinity => visitNonNullString("Infinity", -1)
       case Double.NegativeInfinity => visitNonNullString("-Infinity", -1)
@@ -137,20 +146,20 @@ class BaseCharRenderer[T <: upickle.core.CharOps.Output]
   }
 
 
-  def visitString(s: CharSequence, index: Int) = {
+  def visitString(s: CharSequence, index: Int): T = {
 
     if (s eq null) visitNull(index)
     else visitNonNullString(s, index)
   }
 
-  def visitNonNullString(s: CharSequence, index: Int) = {
+  private def visitNonNullString(s: CharSequence, index: Int) = {
     flushBuffer()
-    upickle.core.RenderUtils.escapeChar(unicodeCharBuilder, elemBuilder, s, escapeUnicode)
+    upickle.core.RenderUtils.escapeChar(null, elemBuilder, s, escapeUnicode)
     flushCharBuilder()
     out
   }
 
-  final def renderIndent() = {
+  final def renderIndent(): Unit = {
     if (indent == -1) ()
     else {
       var i = indent * depth

sjsonnet/src/sjsonnet/Std.scala

@@ -1121,36 +1121,54 @@ class Std {
     builtin(ManifestTomlEx),
     builtinWithDefaults("manifestYamlDoc",
                         "v" -> null,
-                        "indent_array_in_object" -> Val.False(dummyPos)){ (args, pos, ev) =>
+                        "indent_array_in_object" -> Val.False(dummyPos),
+                        "quote_keys" -> Val.True(dummyPos)){ (args, pos, ev) =>
       val v = args(0)
       val indentArrayInObject = args(1)  match {
           case Val.False(_) => false
           case Val.True(_) => true
           case _ => Error.fail("indent_array_in_object has to be a boolean, got" + v.getClass)
         }
+      val quoteKeys = args(2) match {
+        case Val.False(_) => false
+        case Val.True(_) => true
+        case _ => Error.fail("quote_keys has to be a boolean, got " + v.getClass)
+      }
       Materializer.apply0(
         v,
-        new YamlRenderer(indentArrayInObject = indentArrayInObject)
+        new YamlRenderer(indentArrayInObject = indentArrayInObject, quoteKeys = quoteKeys)
       )(ev).toString
     },
     builtinWithDefaults("manifestYamlStream",
                         "v" -> null,
-                        "indent_array_in_object" -> Val.False(dummyPos)){ (args, pos, ev) =>
+                        "indent_array_in_object" -> Val.False(dummyPos),
+                        "c_document_end" -> Val.True(dummyPos),
+                        "quote_keys" -> Val.True(dummyPos)){ (args, _, ev) =>
       val v = args(0)
       val indentArrayInObject = args(1)  match {
         case Val.False(_) => false
         case Val.True(_) => true
         case _ => Error.fail("indent_array_in_object has to be a boolean, got" + v.getClass)
       }
+      val cDocumentEnd = args(2) match {
+        case Val.False(_) => false
+        case Val.True(_) => true
+        case _ => Error.fail("c_document_end has to be a boolean, got " + v.getClass)
+      }
+      val quoteKeys = args(3) match {
+        case Val.False(_) => false
+        case Val.True(_) => true
+        case _ => Error.fail("quote_keys has to be a boolean, got " + v.getClass)
+      }
       v match {
         case arr: Val.Arr => arr.asLazyArray
           .map { item =>
             Materializer.apply0(
               item.force,
-              new YamlRenderer(indentArrayInObject = indentArrayInObject)
+              new YamlRenderer(indentArrayInObject = indentArrayInObject, quoteKeys = quoteKeys)
             )(ev).toString()
           }
-          .mkString("---\n", "\n---\n", "\n...\n")
+          .mkString("---\n", "\n---\n", if (cDocumentEnd) "\n...\n" else "\n")
         case _ => Error.fail("manifestYamlStream only takes arrays, got " + v.getClass)
       }
     },

sjsonnet/src/sjsonnet/YamlRenderer.scala

@@ -1,26 +1,40 @@
 package sjsonnet
 
-import java.io.{StringWriter, Writer}
+import java.io.StringWriter
 import java.util.regex.Pattern
+import upickle.core.{ArrVisitor, ObjVisitor, SimpleVisitor, Visitor}
 
-import upickle.core.{ArrVisitor, ObjVisitor}
+import scala.util.Try
 
 
 
 class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayInObject: Boolean = false,
-                   indent: Int = 2) extends BaseCharRenderer(_out, indent){
+                   quoteKeys: Boolean = true, indent: Int = 2) extends BaseCharRenderer(_out, indent){
   var newlineBuffered = false
   var dashBuffered = false
   var afterKey = false
   private var topLevel = true
+  private val outBuffer = _out.getBuffer
+
+  private val yamlKeyVisitor = new SimpleVisitor[StringWriter, StringWriter]() {
+    override def expectedMsg = "Expected a string key"
+    override def visitString(s: CharSequence, index: Int): StringWriter = {
+      YamlRenderer.this.flushBuffer()
+      if (quoteKeys || !YamlRenderer.isSafeBareKey(s.toString)) {
+        upickle.core.RenderUtils.escapeChar(null, YamlRenderer.this.elemBuilder, s, unicode = true)
+      } else {
+        YamlRenderer.this.appendString(s.toString)
+      }
+      YamlRenderer.this.flushCharBuilder()
+      _out
+    }
+  }
 
-  private val outBuffer = _out.getBuffer()
-
-  override def flushCharBuilder() = {
+  override def flushCharBuilder(): Unit = {
     elemBuilder.writeOutToIfLongerThan(_out, if (depth <= 0 || topLevel) 0 else 1000)
   }
 
-  private[this] def appendString(s: String) = {
+  private[this] def appendString(s: String): Unit = {
     val len = s.length
     var i = 0
     elemBuilder.ensureLength(len)
@@ -48,20 +62,20 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
       }
       depth -= 1
     } else {
-      upickle.core.RenderUtils.escapeChar(unicodeCharBuilder, elemBuilder, s, true)
+      upickle.core.RenderUtils.escapeChar(null, elemBuilder, s, unicode=true)
     }
     flushCharBuilder()
     _out
   }
 
-  override def visitFloat64(d: Double, index: Int) = {
+  override def visitFloat64(d: Double, index: Int): StringWriter = {
     flushBuffer()
     appendString(RenderUtils.renderDouble(d))
     flushCharBuilder()
     _out
   }
 
-  override def flushBuffer() = {
+  override def flushBuffer(): Unit = {
     if (newlineBuffered) {
       // drop space between colon and newline
       elemBuilder.writeOutToIfLongerThan(_out, 0)
@@ -81,7 +95,7 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
     dashBuffered = false
   }
 
-  override def visitArray(length: Int, index: Int) = new ArrVisitor[StringWriter, StringWriter] {
+  override def visitArray(length: Int, index: Int): ArrVisitor[StringWriter, StringWriter] = new ArrVisitor[StringWriter, StringWriter] {
     var empty = true
     flushBuffer()
 
@@ -91,19 +105,19 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
     }
     topLevel = false
 
-    val dedentInObject = afterKey && !indentArrayInObject
+    private val dedentInObject = afterKey && !indentArrayInObject
     afterKey = false
     if (dedentInObject) depth -= 1
     dashBuffered = true
 
-    def subVisitor = YamlRenderer.this
+    def subVisitor: Visitor[StringWriter, StringWriter] = YamlRenderer.this
     def visitValue(v: StringWriter, index: Int): Unit = {
       empty = false
       flushBuffer()
       newlineBuffered = true
       dashBuffered = true
     }
-    def visitEnd(index: Int) = {
+    def visitEnd(index: Int): StringWriter = {
       if (!dedentInObject) depth -= 1
       if (empty) {
         elemBuilder.ensureLength(2)
@@ -116,16 +130,19 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
       _out
     }
   }
-  override def visitObject(length: Int, index: Int) = new ObjVisitor[StringWriter, StringWriter] {
+
+  override def visitObject(length: Int, index: Int): ObjVisitor[StringWriter, StringWriter] = new ObjVisitor[StringWriter, StringWriter] {
     var empty = true
     flushBuffer()
     if (!topLevel) depth += 1
     topLevel = false
 
     if (afterKey) newlineBuffered = true
 
-    def subVisitor = YamlRenderer.this
-    def visitKey(index: Int) = YamlRenderer.this
+    def subVisitor: Visitor[StringWriter, StringWriter] = YamlRenderer.this
+
+    def visitKey(index: Int): Visitor[StringWriter, StringWriter] = yamlKeyVisitor
+
     def visitKeyValue(s: Any): Unit = {
       empty = false
       flushBuffer()
@@ -136,11 +153,13 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
       afterKey = true
       newlineBuffered = false
     }
+
     def visitValue(v: StringWriter, index: Int): Unit = {
       newlineBuffered = true
       afterKey = false
     }
-    def visitEnd(index: Int) = {
+
+    def visitEnd(index: Int): StringWriter = {
       if (empty) {
         elemBuilder.ensureLength(2)
         elemBuilder.append('{')
@@ -155,9 +174,28 @@ class YamlRenderer(_out: StringWriter = new java.io.StringWriter(), indentArrayI
   }
 }
 object YamlRenderer{
-  val newlinePattern = Pattern.compile("\n")
+  val newlinePattern: Pattern = Pattern.compile("\n")
+  private val safeYamlKeyPattern = Pattern.compile("^[a-zA-Z0-9/._-]+$")
+  private val yamlReserved = Set("true", "false", "null", "yes", "no", "on", "off", "y", "n", ".nan",
+    "+.inf", "-.inf", ".inf", "null", "-", "---", "''")
+  private val yamlTimestampPattern = Pattern.compile("^(?:[0-9]*-){2}[0-9]*$")
+  private val yamlBinaryPattern = Pattern.compile("^[-+]?0b[0-1_]+$")
+  private val yamlHexPattern = Pattern.compile("[-+]?0x[0-9a-fA-F_]+")
+  private val yamlFloatPattern = Pattern.compile( "^-?([0-9_]*)*(\\.[0-9_]*)?(e[-+][0-9_]+)?$" )
+  private val yamlIntPattern = Pattern.compile("^[-+]?[0-9_]+$")
+
+  private def isSafeBareKey(k: String) = {
+    val l = k.toLowerCase
+    !yamlReserved.contains(l) &&
+      safeYamlKeyPattern.matcher(k).matches() &&
+      !yamlTimestampPattern.matcher(l).matches() &&
+      !yamlBinaryPattern.matcher(k).matches() &&
+      !yamlHexPattern.matcher(k).matches() &&
+      !yamlFloatPattern.matcher(l).matches() &&
+      !yamlIntPattern.matcher(l).matches()
+  }
 
-  def writeIndentation(out: upickle.core.CharBuilder, n: Int) = {
+  def writeIndentation(out: upickle.core.CharBuilder, n: Int): Unit = {
     out.ensureLength(n+1)
     out.append('\n')
     var i = n