diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatter.scala index 2087cc7f109bc1bf712f2b17cc77176d45ad2581..c98182c96b165848126cf2d8183b5ca5c759bab4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatter.scala @@ -18,8 +18,7 @@ package org.apache.spark.sql.catalyst.expressions.codegen /** - * An utility class that indents a block of code based on the curly braces. - * + * An utility class that indents a block of code based on the curly braces and parentheses. * This is used to prettify generated code when in debug mode (or exceptions). * * Written by Matei Zaharia. @@ -35,11 +34,12 @@ private class CodeFormatter { private var indentString = "" private def addLine(line: String): Unit = { - val indentChange = line.count(_ == '{') - line.count(_ == '}') + val indentChange = + line.count(c => "({".indexOf(c) >= 0) - line.count(c => ")}".indexOf(c) >= 0) val newIndentLevel = math.max(0, indentLevel + indentChange) // Lines starting with '}' should be de-indented even if they contain '{' after; // in addition, lines ending with ':' are typically labels - val thisLineIndent = if (line.startsWith("}") || line.endsWith(":")) { + val thisLineIndent = if (line.startsWith("}") || line.startsWith(")") || line.endsWith(":")) { " " * (indentSize * (indentLevel - 1)) } else { indentString diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatterSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatterSuite.scala index 478702fea61460604a7246b1ae3390300facce6e..46daa3eb8bf800bbe13a1121e5d5f1fe593e5c8b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatterSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeFormatterSuite.scala @@ -73,4 +73,34 @@ class CodeFormatterSuite extends SparkFunSuite { |} """.stripMargin } + + testCase("if else on the same line") { + """ + |class A { + | if (c) {duh;} else {boo;} + |} + """.stripMargin + }{ + """ + |class A { + | if (c) {duh;} else {boo;} + |} + """.stripMargin + } + + testCase("function calls") { + """ + |foo( + |a, + |b, + |c) + """.stripMargin + }{ + """ + |foo( + | a, + | b, + | c) + """.stripMargin + } }