Fix exec mode - add \n on end if not present

ncw · ncw · commit 378d65e34ec1 · 2014-11-10T21:07:33.000Z
diff --git a/parser/lexer.go b/parser/lexer.go
@@ -41,7 +41,8 @@ type yyLex struct {
 	indentStack   []int   // indent stack to control INDENT / DEDENT tokens
 	state         int     // current state of state machine
 	currentIndent string  // whitespace at start of current line
-	interactive   bool    // set if reading interactive input
+	interactive   bool    // set if mode "single" reading interactive input
+	exec          bool    // set if mode "exec" reading from file
 	bracket       int     // number of open [ ]
 	parenthesis   int     // number of open ( )
 	brace         int     // number of open { }
@@ -64,6 +65,7 @@ func NewLex(r io.Reader, mode string) (*yyLex, error) {
 	switch mode {
 	case "exec":
 		x.startToken = FILE_INPUT
+		x.exec = true
 	case "eval":
 		x.startToken = EVAL_INPUT
 	case "single":
@@ -90,6 +92,11 @@ func (x *yyLex) refill() {
 		x.eof = true
 		x.Errorf("Error reading input: %v", err)
 	}
+	// If this is exec input, add a newline to the end of the
+	// string if there isn't one already.
+	if x.eof && x.exec && len(x.line) > 0 && x.line[len(x.line)-1] != '\n' {
+		x.line += "\n"
+	}
 }
 
 // Finds the length of a space and tab seperated string
diff --git a/parser/lexer_test.go b/parser/lexer_test.go
@@ -263,24 +263,25 @@ func TestLex(t *testing.T) {
 			{'(', nil},
 			{NUMBER, py.Int(1)},
 			{')', nil},
+			{NEWLINE, nil},
 			{ENDMARKER, nil},
 		}},
-		{"{\n  1\n}", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"{\n  1\n}", "", "single", LexTokens{
+			{SINGLE_INPUT, nil},
 			{'{', nil},
 			{NUMBER, py.Int(1)},
 			{'}', nil},
 			{ENDMARKER, nil},
 		}},
-		{"[\n  1\n]", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"[\n  1\n]", "", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 			{'[', nil},
 			{NUMBER, py.Int(1)},
 			{']', nil},
 			{ENDMARKER, nil},
 		}},
-		{"1\\\n2", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"1\\\n2", "", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 			{NUMBER, py.Int(1)},
 			{NUMBER, py.Int(2)},
 			{ENDMARKER, nil},
@@ -295,33 +296,33 @@ func TestLex(t *testing.T) {
 			{NUMBER, py.Int(1)},
 			{ENDMARKER, nil},
 		}},
-		{"'1\\\n2'", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"'1\\\n2'", "", "single", LexTokens{
+			{SINGLE_INPUT, nil},
 			{STRING, py.String("12")},
 			{ENDMARKER, nil},
 		}},
-		{"0x1234 +\t0.1-6.1j", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"0x1234 +\t0.1-6.1j", "", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 			{NUMBER, py.Int(0x1234)},
 			{'+', nil},
 			{NUMBER, py.Float(0.1)},
 			{'-', nil},
 			{NUMBER, py.Complex(complex(0, 6.1))},
 			{ENDMARKER, nil},
 		}},
-		{"001", "illegal decimal with leading zero", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"001", "illegal decimal with leading zero", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 		}},
-		{"u'''1\n2\n'''", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"u'''1\n2\n'''", "", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 			{STRING, py.String("1\n2\n")},
 			{ENDMARKER, nil},
 		}},
-		{"\"hello\n", "Unterminated \"x\" string", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"\"hello\n", "Unterminated \"x\" string", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 		}},
-		{"1 >>-3\na <<=+12", "", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"1 >>-3\na <<=+12", "", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 			{NUMBER, py.Int(1)},
 			{GTGT, nil},
 			{'-', nil},
@@ -333,8 +334,8 @@ func TestLex(t *testing.T) {
 			{NUMBER, py.Int(12)},
 			{ENDMARKER, nil},
 		}},
-		{"$asdasd", "invalid syntax", "exec", LexTokens{
-			{FILE_INPUT, nil},
+		{"$asdasd", "invalid syntax", "eval", LexTokens{
+			{EVAL_INPUT, nil},
 		}},
 	} {
 		lts, err := LexString(test.in, test.mode)
@@ -595,7 +596,7 @@ func TestLexerReadString(t *testing.T) {
 		{`BR"""a\nc"""`, STRING, py.Bytes(string(`a\nc`)), ``},
 		{`rB'''a\"c'''`, STRING, py.Bytes(string(`a\"c`)), ``},
 	} {
-		x, err := NewLex(bytes.NewBufferString(test.in), "exec")
+		x, err := NewLex(bytes.NewBufferString(test.in), "eval")
 		if err != nil {
 			t.Fatal(err)
 		}
