pretty print improvements (apache#1851)

Author: lovasoa

README.md

@@ -89,10 +89,14 @@ keywords, the following should hold true for all SQL:
 
 ```rust
 // Parse SQL
+let sql = "SELECT 'hello'";
 let ast = Parser::parse_sql(&GenericDialect, sql).unwrap();
 
 // The original SQL text can be generated from the AST
 assert_eq!(ast[0].to_string(), sql);
+
+// The SQL can also be pretty-printed with newlines and indentation
+assert_eq!(format!("{:#}", ast[0]), "SELECT\n  'hello'");
 ```
 
 There are still some cases in this crate where different SQL with seemingly

src/ast/dml.rs

@@ -29,6 +29,8 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
+use crate::display_utils::{indented_list, Indent, SpaceOrNewline};
+
 pub use super::ddl::{ColumnDef, TableConstraint};
 
 use super::{
@@ -579,28 +581,32 @@ impl Display for Insert {
             )?;
         }
         if !self.columns.is_empty() {
-            write!(f, "({}) ", display_comma_separated(&self.columns))?;
+            write!(f, "({})", display_comma_separated(&self.columns))?;
+            SpaceOrNewline.fmt(f)?;
         }
         if let Some(ref parts) = self.partitioned {
             if !parts.is_empty() {
-                write!(f, "PARTITION ({}) ", display_comma_separated(parts))?;
+                write!(f, "PARTITION ({})", display_comma_separated(parts))?;
+                SpaceOrNewline.fmt(f)?;
             }
         }
         if !self.after_columns.is_empty() {
-            write!(f, "({}) ", display_comma_separated(&self.after_columns))?;
+            write!(f, "({})", display_comma_separated(&self.after_columns))?;
+            SpaceOrNewline.fmt(f)?;
         }
 
         if let Some(settings) = &self.settings {
-            write!(f, "SETTINGS {} ", display_comma_separated(settings))?;
+            write!(f, "SETTINGS {}", display_comma_separated(settings))?;
+            SpaceOrNewline.fmt(f)?;
         }
 
         if let Some(source) = &self.source {
-            write!(f, "{source}")?;
+            source.fmt(f)?;
         } else if !self.assignments.is_empty() {
-            write!(f, "SET ")?;
-            write!(f, "{}", display_comma_separated(&self.assignments))?;
+            write!(f, "SET")?;
+            indented_list(f, &self.assignments)?;
         } else if let Some(format_clause) = &self.format_clause {
-            write!(f, "{format_clause}")?;
+            format_clause.fmt(f)?;
         } else if self.columns.is_empty() {
             write!(f, "DEFAULT VALUES")?;
         }
@@ -620,7 +626,9 @@ impl Display for Insert {
         }
 
         if let Some(returning) = &self.returning {
-            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("RETURNING")?;
+            indented_list(f, returning)?;
         }
         Ok(())
     }
@@ -649,32 +657,45 @@ pub struct Delete {
 
 impl Display for Delete {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "DELETE ")?;
+        f.write_str("DELETE")?;
         if !self.tables.is_empty() {
-            write!(f, "{} ", display_comma_separated(&self.tables))?;
+            indented_list(f, &self.tables)?;
         }
         match &self.from {
             FromTable::WithFromKeyword(from) => {
-                write!(f, "FROM {}", display_comma_separated(from))?;
+                f.write_str(" FROM")?;
+                indented_list(f, from)?;
             }
             FromTable::WithoutKeyword(from) => {
-                write!(f, "{}", display_comma_separated(from))?;
+                indented_list(f, from)?;
             }
         }
         if let Some(using) = &self.using {
-            write!(f, " USING {}", display_comma_separated(using))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("USING")?;
+            indented_list(f, using)?;
         }
         if let Some(selection) = &self.selection {
-            write!(f, " WHERE {selection}")?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("WHERE")?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(selection).fmt(f)?;
         }
         if let Some(returning) = &self.returning {
-            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("RETURNING")?;
+            indented_list(f, returning)?;
         }
         if !self.order_by.is_empty() {
-            write!(f, " ORDER BY {}", display_comma_separated(&self.order_by))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("ORDER BY")?;
+            indented_list(f, &self.order_by)?;
         }
         if let Some(limit) = &self.limit {
-            write!(f, " LIMIT {limit}")?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("LIMIT")?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(limit).fmt(f)?;
         }
         Ok(())
     }

src/ast/mod.rs

@@ -41,7 +41,7 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::{
-    display_utils::SpaceOrNewline,
+    display_utils::{indented_list, SpaceOrNewline},
     tokenizer::{Span, Token},
 };
 use crate::{
@@ -4548,7 +4548,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Insert(insert) => write!(f, "{insert}"),
+            Statement::Insert(insert) => insert.fmt(f),
             Statement::Install {
                 extension_name: name,
             } => write!(f, "INSTALL {name}"),
@@ -4611,30 +4611,42 @@ impl fmt::Display for Statement {
                 returning,
                 or,
             } => {
-                write!(f, "UPDATE ")?;
+                f.write_str("UPDATE ")?;
                 if let Some(or) = or {
-                    write!(f, "{or} ")?;
+                    or.fmt(f)?;
+                    f.write_str(" ")?;
                 }
-                write!(f, "{table}")?;
+                table.fmt(f)?;
                 if let Some(UpdateTableFromKind::BeforeSet(from)) = from {
-                    write!(f, " FROM {}", display_comma_separated(from))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("FROM")?;
+                    indented_list(f, from)?;
                 }
                 if !assignments.is_empty() {
-                    write!(f, " SET {}", display_comma_separated(assignments))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("SET")?;
+                    indented_list(f, assignments)?;
                 }
                 if let Some(UpdateTableFromKind::AfterSet(from)) = from {
-                    write!(f, " FROM {}", display_comma_separated(from))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("FROM")?;
+                    indented_list(f, from)?;
                 }
                 if let Some(selection) = selection {
-                    write!(f, " WHERE {selection}")?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("WHERE")?;
+                    SpaceOrNewline.fmt(f)?;
+                    Indent(selection).fmt(f)?;
                 }
                 if let Some(returning) = returning {
-                    write!(f, " RETURNING {}", display_comma_separated(returning))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("RETURNING")?;
+                    indented_list(f, returning)?;
                 }
                 Ok(())
             }
-            Statement::Delete(delete) => write!(f, "{delete}"),
-            Statement::Open(open) => write!(f, "{open}"),
+            Statement::Delete(delete) => delete.fmt(f),
+            Statement::Open(open) => open.fmt(f),
             Statement::Close { cursor } => {
                 write!(f, "CLOSE {cursor}")?;
 

src/ast/query.rs

@@ -2888,13 +2888,14 @@ pub struct Values {
 
 impl fmt::Display for Values {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "VALUES ")?;
+        f.write_str("VALUES")?;
         let prefix = if self.explicit_row { "ROW" } else { "" };
         let mut delim = "";
         for row in &self.rows {
-            write!(f, "{delim}")?;
-            delim = ", ";
-            write!(f, "{prefix}({})", display_comma_separated(row))?;
+            f.write_str(delim)?;
+            delim = ",";
+            SpaceOrNewline.fmt(f)?;
+            Indent(format_args!("{prefix}({})", display_comma_separated(row))).fmt(f)?;
         }
         Ok(())
     }

src/display_utils.rs

@@ -31,13 +31,10 @@ where
     T: Write,
 {
     fn write_str(&mut self, s: &str) -> fmt::Result {
-        let mut first = true;
-        for line in s.split('\n') {
-            if !first {
-                write!(self.0, "\n{INDENT}")?;
-            }
-            self.0.write_str(line)?;
-            first = false;
+        self.0.write_str(s)?;
+        // Our NewLine and SpaceOrNewline utils always print individual newlines as a single-character string.
+        if s == "\n" {
+            self.0.write_str(INDENT)?;
         }
         Ok(())
     }
@@ -71,7 +68,7 @@ impl Display for SpaceOrNewline {
 
 /// A value that displays a comma-separated list of values.
 /// When pretty-printed (using {:#}), it displays each value on a new line.
-pub struct DisplayCommaSeparated<'a, T: fmt::Display>(&'a [T]);
+pub(crate) struct DisplayCommaSeparated<'a, T: fmt::Display>(&'a [T]);
 
 impl<T: fmt::Display> fmt::Display for DisplayCommaSeparated<'_, T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -89,45 +86,33 @@ impl<T: fmt::Display> fmt::Display for DisplayCommaSeparated<'_, T> {
 }
 
 /// Displays a whitespace, followed by a comma-separated list that is indented when pretty-printed.
-pub(crate) fn indented_list<T: fmt::Display>(f: &mut fmt::Formatter, slice: &[T]) -> fmt::Result {
+pub(crate) fn indented_list<T: fmt::Display>(f: &mut fmt::Formatter, items: &[T]) -> fmt::Result {
     SpaceOrNewline.fmt(f)?;
-    Indent(DisplayCommaSeparated(slice)).fmt(f)
+    Indent(DisplayCommaSeparated(items)).fmt(f)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    struct DisplayCharByChar<T: Display>(T);
+    #[test]
+    fn test_indent() {
+        struct TwoLines;
 
-    impl<T: Display> Display for DisplayCharByChar<T> {
-        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-            for c in self.0.to_string().chars() {
-                write!(f, "{}", c)?;
+        impl Display for TwoLines {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.write_str("line 1")?;
+                SpaceOrNewline.fmt(f)?;
+                f.write_str("line 2")
             }
-            Ok(())
         }
-    }
 
-    #[test]
-    fn test_indent() {
-        let original = "line 1\nline 2";
-        let indent = Indent(original);
+        let indent = Indent(TwoLines);
         assert_eq!(
             indent.to_string(),
-            original,
+            TwoLines.to_string(),
             "Only the alternate form should be indented"
         );
-        let expected = "  line 1\n  line 2";
-        assert_eq!(format!("{:#}", indent), expected);
-        let display_char_by_char = DisplayCharByChar(original);
-        assert_eq!(format!("{:#}", Indent(display_char_by_char)), expected);
-    }
-
-    #[test]
-    fn test_space_or_newline() {
-        let space_or_newline = SpaceOrNewline;
-        assert_eq!(format!("{}", space_or_newline), " ");
-        assert_eq!(format!("{:#}", space_or_newline), "\n");
+        assert_eq!(format!("{:#}", indent), "  line 1\n  line 2");
     }
 }