lexer: fix whitespace issues

ncw · ncw · commit 6ab706154053 · 2015-03-27T19:45:10.000Z
diff --git a/parser/grammar_data_test.go b/parser/grammar_data_test.go
@@ -203,6 +203,7 @@ var grammarTestData = []struct {
 	{"while True:\n pass\nelse:\n return\n", "exec", "Module(body=[While(test=NameConstant(value=True), body=[Pass()], orelse=[Return(value=None)])])"},
 	{"if True: pass", "exec", "Module(body=[If(test=NameConstant(value=True), body=[Pass()], orelse=[])])"},
 	{"if True:\n pass\n", "exec", "Module(body=[If(test=NameConstant(value=True), body=[Pass()], orelse=[])])"},
+	{"if True:\n pass\n\n", "exec", "Module(body=[If(test=NameConstant(value=True), body=[Pass()], orelse=[])])"},
 	{"if True:\n    pass\n    continue\nelse:\n    break\n    pass\n", "exec", "Module(body=[If(test=NameConstant(value=True), body=[Pass(), Continue()], orelse=[Break(), Pass()])])"},
 	{"if a:\n    continue\nelif b:\n    break\nelif c:\n    pass\nelif c:\n    continue\n    pass\n", "exec", "Module(body=[If(test=Name(id='a', ctx=Load()), body=[Continue()], orelse=[If(test=Name(id='b', ctx=Load()), body=[Break()], orelse=[If(test=Name(id='c', ctx=Load()), body=[Pass()], orelse=[If(test=Name(id='c', ctx=Load()), body=[Continue(), Pass()], orelse=[])])])])])"},
 	{"if a:\n    continue\nelif b:\n    break\nelse:\n    continue\n    pass\n", "exec", "Module(body=[If(test=Name(id='a', ctx=Load()), body=[Continue()], orelse=[If(test=Name(id='b', ctx=Load()), body=[Break()], orelse=[Continue(), Pass()])])])"},
diff --git a/parser/lexer.go b/parser/lexer.go
@@ -372,12 +372,10 @@ func (x *yyLex) Lex(yylval *yySymType) (ret int) {
 			x.line = trimmed
 			x.state++
 		case checkEmpty:
+			despaced := strings.TrimSpace(x.line) // remove other whitespace other than " \t"
 			// Ignore line if just white space or whitespace then comment
-			if x.line == "" || x.line == "\n" || x.line[0] == '#' {
+			if despaced == "" || despaced[0] == '#' {
 				x.state = checkEof
-				if x.line != "" {
-					return NEWLINE
-				}
 				continue
 			}
 			x.state++
diff --git a/parser/lexer_test.go b/parser/lexer_test.go
@@ -204,7 +204,6 @@ func TestLex(t *testing.T) {
 		{"\n", "", "single", LexTokens{
 			{SINGLE_INPUT, nil},
 			{NEWLINE, nil},
-			{NEWLINE, nil},
 		}},
 		{"pass", "", "single", LexTokens{
 			{SINGLE_INPUT, nil},
@@ -218,9 +217,10 @@ func TestLex(t *testing.T) {
 		}},
 		{"\n#hello\n  #comment\n", "", "exec", LexTokens{
 			{FILE_INPUT, nil},
-			{NEWLINE, nil},
-			{NEWLINE, nil},
-			{NEWLINE, nil},
+			{ENDMARKER, nil},
+		}},
+		{"\n#hello\n\f  #comment\n", "", "exec", LexTokens{
+			{FILE_INPUT, nil},
 			{ENDMARKER, nil},
 		}},
 		{"1\n 2\n", "", "eval", LexTokens{
diff --git a/parser/make_grammar_test.py b/parser/make_grammar_test.py
@@ -232,6 +232,7 @@
     ("while True:\n pass\nelse:\n return\n", "exec"),
     ("if True: pass", "exec"),
     ("if True:\n pass\n", "exec"),
+    ("if True:\n pass\n\n", "exec"),
     ("""\
 if True:
     pass
