Fix formatting of lib's fun starting with uppercase letter

chris-huxtable · Jan 25, 2018 · b793876 · b793876
1 parent 0ebc173
commit b793876
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/spec/compiler/formatter/formatter_spec.cr b/spec/compiler/formatter/formatter_spec.cr
@@ -581,6 +581,9 @@ describe Crystal::Formatter do
   assert_format "lib Foo\nend"
   assert_format "lib Foo\ntype  Foo  =   Bar\nend", "lib Foo\n  type Foo = Bar\nend"
   assert_format "lib Foo\nfun foo\nend", "lib Foo\n  fun foo\nend"
+  assert_format "lib Foo\n  fun Bar\nend"
+  assert_format "lib Foo\n  fun bar = Bar\nend"
+  assert_format "lib Foo\n  fun Foo = Bar\nend"
   assert_format "lib Foo\nfun foo  :  Int32\nend", "lib Foo\n  fun foo : Int32\nend"
   assert_format "lib Foo\nfun foo()  :  Int32\nend", "lib Foo\n  fun foo : Int32\nend"
   assert_format "lib Foo\nfun foo ()  :  Int32\nend", "lib Foo\n  fun foo : Int32\nend"

diff --git a/src/compiler/crystal/tools/formatter.cr b/src/compiler/crystal/tools/formatter.cr
@@ -1530,7 +1530,7 @@ module Crystal
     def visit(node : FunDef)
       write_keyword :fun, " "
 
-      check :IDENT
+      check :IDENT, :CONST
       write node.name
       next_token_skip_space
 
@@ -4573,8 +4573,10 @@ module Crystal
       raise "expecting keyword #{keywords.join " or "}, not `#{@token.type}, #{@token.value}`, at #{@token.location}" unless keywords.any? { |k| @token.keyword?(k) }
     end
 
-    def check(token_type)
-      raise "expecting #{token_type}, not `#{@token.type}, #{@token.value}`, at #{@token.location}" unless @token.type == token_type
+    def check(*token_types)
+      unless token_types.includes? @token.type
+        raise "expecting #{token_types.join " or "}, not `#{@token.type}, #{@token.value}`, at #{@token.location}"
+      end
     end
 
     def check_end