Merge pull request #267 from hhatto/feature-e731

fix e731

autopep8.py

@@ -75,6 +75,7 @@ class documentation for more information.
 
 
 PYTHON_SHEBANG_REGEX = re.compile(r'^#!.*\bpython[23]?\b\s*$')
+LAMBDA_REGEX = re.compile(r'([\w.]+)\s=\slambda\s+([\w,\s.]+):')
 COMPARE_NEGATIVE_REGEX = re.compile(r'\b(not)\s+([^][)(}{ ]+)\s+(in|is)\s')
 
 
@@ -404,6 +405,7 @@ class FixPEP8(object):
         - e502
         - e701,e702
         - e711,e712,e713,e714
+        - e731
         - w291
 
     """
@@ -999,6 +1001,17 @@ def fix_e714(self, result):
                     target[:pos_start], match.group(2), match.group(3),
                     match.group(1), target[match.end():])
 
+    def fix_e731(self, result):
+        """Fix do not assign a lambda expression check."""
+        (line_index, _, target) = get_index_offset_contents(result,
+                                                            self.source)
+        match = LAMBDA_REGEX.search(target)
+        if match:
+            end = match.end()
+            self.source[line_index] = "{0}def {1}({2}): return {3}".format(
+                target[:match.start(0)], match.group(1), match.group(2),
+                target[end:])
+
     def fix_w291(self, result):
         """Remove trailing whitespace."""
         fixed_line = self.source[result['line'] - 1].rstrip()

test/test_autopep8.py

@@ -3773,6 +3773,18 @@ def test_e721_in_conditional(self):
         with autopep8_context(line, options=['--aggressive']) as result:
             self.assertEqual(fixed, result)
 
+    def test_e731(self):
+        line = 'a = lambda x: x * 2\n'
+        fixed = 'def a(x): return x * 2\n'
+        with autopep8_context(line) as result:
+            self.assertEqual(fixed, result)
+
+    def test_e731_with_args(self):
+        line = 'a = lambda x, y: x * 2 + y\n'
+        fixed = 'def a(x, y): return x * 2 + y\n'
+        with autopep8_context(line) as result:
+            self.assertEqual(fixed, result)
+
     def test_e901_should_cause_indentation_screw_up(self):
         line = """\
 def tmp(g):