From b3e9653b7fb4c78ccd5f7c23e402f8a3f678c924 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Mon, 2 Feb 2015 22:54:46 -0800
Subject: [PATCH] Properly escape keys when printing TOML

Closes #53
---
 src/display.rs | 77 +++++++++++++++++++++++++++++++++++---------------
 tests/valid.rs | 10 +++++--
 2 files changed, 62 insertions(+), 25 deletions(-)

diff --git a/src/display.rs b/src/display.rs
index 21aec1e..f1b2a76 100644
--- a/src/display.rs
+++ b/src/display.rs
@@ -8,25 +8,12 @@ struct Printer<'a, 'b:'a> {
     stack: Vec<&'a str>,
 }
 
+struct Key<'a>(&'a [&'a str]);
+
 impl fmt::Display for Value {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
-            String(ref s) => {
-                try!(write!(f, "\""));
-                for ch in s.chars() {
-                    match ch {
-                        '\u{8}' => try!(write!(f, "\\b")),
-                        '\u{9}' => try!(write!(f, "\\t")),
-                        '\u{a}' => try!(write!(f, "\\n")),
-                        '\u{c}' => try!(write!(f, "\\f")),
-                        '\u{d}' => try!(write!(f, "\\r")),
-                        '\u{22}' => try!(write!(f, "\\\"")),
-                        '\u{5c}' => try!(write!(f, "\\\\")),
-                        ch => try!(write!(f, "{}", ch)),
-                    }
-                }
-                write!(f, "\"")
-            }
+            String(ref s) => write_str(f, s),
             Integer(i) => write!(f, "{}", i),
             Float(fp) => {
                 try!(write!(f, "{}", fp));
@@ -51,6 +38,23 @@ impl fmt::Display for Value {
     }
 }
 
+fn write_str(f: &mut fmt::Formatter, s: &str) -> fmt::Result {
+    try!(write!(f, "\""));
+    for ch in s.chars() {
+        match ch {
+            '\u{8}' => try!(write!(f, "\\b")),
+            '\u{9}' => try!(write!(f, "\\t")),
+            '\u{a}' => try!(write!(f, "\\n")),
+            '\u{c}' => try!(write!(f, "\\f")),
+            '\u{d}' => try!(write!(f, "\\r")),
+            '\u{22}' => try!(write!(f, "\\\"")),
+            '\u{5c}' => try!(write!(f, "\\\\")),
+            ch => try!(write!(f, "{}", ch)),
+        }
+    }
+    write!(f, "\"")
+}
+
 impl<'a, 'b> Printer<'a, 'b> {
     fn print(&mut self, table: &'a TomlTable) -> fmt::Result {
         for (k, v) in table.iter() {
@@ -64,14 +68,13 @@ impl<'a, 'b> Printer<'a, 'b> {
                 }
                 _ => {}
             }
-            try!(writeln!(self.output, "{} = {}", k, v));
+            try!(writeln!(self.output, "{} = {}", Key(&[k]), v));
         }
         for (k, v) in table.iter() {
             match *v {
                 Table(ref inner) => {
-                    self.stack.push(&**k);
-                    try!(writeln!(self.output, "\n[{}]",
-                                  self.stack.connect(".")));
+                    self.stack.push(k);
+                    try!(writeln!(self.output, "\n[{}]", Key(&self.stack)));
                     try!(self.print(inner));
                     self.stack.pop();
                 }
@@ -80,10 +83,9 @@ impl<'a, 'b> Printer<'a, 'b> {
                         Some(&Table(..)) => {}
                         _ => continue
                     }
-                    self.stack.push(&**k);
+                    self.stack.push(k);
                     for inner in inner.iter() {
-                        try!(writeln!(self.output, "\n[[{}]]",
-                                      self.stack.connect(".")));
+                        try!(writeln!(self.output, "\n[[{}]]", Key(&self.stack)));
                         match *inner {
                             Table(ref inner) => try!(self.print(inner)),
                             _ => panic!("non-heterogeneous toml array"),
@@ -98,6 +100,29 @@ impl<'a, 'b> Printer<'a, 'b> {
     }
 }
 
+impl<'a> fmt::Display for Key<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        for (i, part) in self.0.iter().enumerate() {
+            if i != 0 { try!(write!(f, ".")); }
+            let ok = part.chars().all(|c| {
+                match c {
+                    'a' ... 'z' |
+                    'A' ... 'Z' |
+                    '0' ... '9' |
+                    '-' | '_' => true,
+                    _ => false,
+                }
+            });
+            if ok {
+                try!(write!(f, "{}", part));
+            } else {
+                try!(write_str(f, part));
+            }
+        }
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 #[allow(warnings)]
 mod tests {
@@ -167,5 +192,11 @@ mod tests {
                     \n\
                     [[test2]]\n\
                     test = \"wut\"\n");
+        assert_eq!(Table(map! {
+                        "foo.bar" => Integer(2),
+                        "foo\"bar" => Integer(2)
+                   }).to_string().as_slice(),
+                   "\"foo\\\"bar\" = 2\n\
+                    \"foo.bar\" = 2\n");
     }
 }
diff --git a/tests/valid.rs b/tests/valid.rs
index 628130e..1e47829 100644
--- a/tests/valid.rs
+++ b/tests/valid.rs
@@ -46,14 +46,20 @@ fn run(toml: &str, json: &str) {
                 (e.desc.clone(), toml.slice(e.lo - 5, e.hi + 5))
             }).collect::<Vec<(String, &str)>>());
     assert!(table.is_some());
-    let table = table.unwrap();
+    let toml = Table(table.unwrap());
+    let toml_string = format!("{}", toml);
 
     let json = Json::from_str(json).unwrap();
-    let toml_json = to_json(Table(table));
+    let toml_json = to_json(toml.clone());
     assert!(json == toml_json,
             "expected\n{}\ngot\n{}\n",
             json.pretty(),
             toml_json.pretty());
+
+    let table2 = Parser::new(&toml_string).parse().unwrap();
+    // floats are a little lossy
+    if table2.values().any(|v| v.as_float().is_some()) { return }
+    assert_eq!(toml, Table(table2));
 }
 
 macro_rules! test( ($name:ident, $toml:expr, $json:expr) => (