Fix argument position when lambdas are used in arguments.

Trundle · Trundle · commit 02d62a9d085e · 2010-11-09T15:02:44.000+01:00
Closes #158.
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -473,11 +473,15 @@ def get_args(self):
                             stack[-1][1] += 1
                         except TypeError:
                             stack[-1][1] = ''
+                    elif value == ':' and stack[-1][2] == 'lambda':
+                        stack.pop()
                 elif (token is Token.Name or token in Token.Name.subtypes or
                       token is Token.Operator and value == '.'):
                     stack[-1][0] += value
                 elif token is Token.Operator and value == '=':
                     stack[-1][1] = stack[-1][0]
+                elif token is Token.Keyword and value == 'lambda':
+                    stack.append(['', 0, value])
                 else:
                     stack[-1][0] = ''
             while stack[-1][2] in '[{':
diff --git a/bpython/test/test_repl.py b/bpython/test/test_repl.py
@@ -1,7 +1,8 @@
+import os
 import unittest
 from itertools import islice
 
-from bpython import repl
+from bpython import config, repl
 
 
 class TestHistory(unittest.TestCase):
@@ -129,6 +130,34 @@ def test_update(self):
         self.assertNotEqual(list(slice), self.matches)
         self.assertEqual(list(newslice), newmatches)
 
+class TestRepl(repl.Repl):
+    def __init__(self):
+        config_struct = config.Struct()
+        config.loadini(config_struct, os.devnull)
+        repl.Repl.__init__(self, repl.Interpreter(), config_struct)
+        self.input_line = ""
+
+    def current_line(self):
+        return self.input_line
+
+class TestArgspec(unittest.TestCase):
+    def setUp(self):
+        self.repl = TestRepl()
+        self.repl.push("def spam(a, b, c):\n", False)
+        self.repl.push("    pass\n", False)
+        self.repl.push("\n", False)
+
+    def setInputLine(self, line):
+        """Set current input line of the test REPL."""
+        self.repl.input_line = line
+
+    def test_lambda_position(self):
+        self.setInputLine("spam(lambda a, b: 1, ")
+        self.assertTrue(self.repl.get_args())
+        self.assertTrue(self.repl.argspec)
+        # Argument position
+        self.assertEqual(self.repl.argspec[3], 1)
+
 
 if __name__ == '__main__':
     unittest.main()
