Bugfix: python add blank lines (#1944)

* Don't add newlines when not a control statement for Python
* Add test for accidental newline fix
* Add docstring detection to avoid adding unnecessarily newlines
* Add tests for docstring detection

2 files changed, 72 insertions, 1 deletions

diff --git a/autoload/ale/fixers/generic_python.vim b/autoload/ale/fixers/generic_python.vim
index 124146be..d55a23c3 100644
--- a/autoload/ale/fixers/generic_python.vim
+++ b/autoload/ale/fixers/generic_python.vim
@@ -6,13 +6,28 @@ function! ale#fixers#generic_python#AddLinesBeforeControlStatements(buffer, line
     let l:new_lines = []
     let l:last_indent_size = 0
     let l:last_line_is_blank = 0
+    let l:in_docstring = 0
 
     for l:line in a:lines
         let l:indent_size = len(matchstr(l:line, '^ *'))
 
+        if !l:in_docstring
+            " Make sure it is not just a single line docstring and then verify
+            " it's starting a new docstring
+            if match(l:line, '\v^ *("""|'''''').*("""|'''''')') == -1
+            \&& match(l:line, '\v^ *("""|'''''')') >= 0
+                let l:in_docstring = 1
+            endif
+        else
+            if match(l:line, '\v^ *.*("""|'''''')') >= 0
+                let l:in_docstring = 0
+            endif
+        endif
+
         if !l:last_line_is_blank
+        \&& !l:in_docstring
         \&& l:indent_size <= l:last_indent_size
-        \&& match(l:line, '\v^ *(return|if|for|while|break|continue)') >= 0
+        \&& match(l:line, '\v^ *(return|if|for|while|break|continue)(\(| |$)') >= 0
             call add(l:new_lines, '')
         endif
 
diff --git a/test/fixers/test_python_add_blank_lines_fixer.vader b/test/fixers/test_python_add_blank_lines_fixer.vader
index 4a91aa10..7d042c8a 100644
--- a/test/fixers/test_python_add_blank_lines_fixer.vader
+++ b/test/fixers/test_python_add_blank_lines_fixer.vader
@@ -6,15 +6,22 @@ After:
 
 Given python(Some Python without blank lines):
   def foo():
+      """ This is a simple test docstring """
       return 1
 
 
   def bar():
+      '''This is another simple test docstring'''
       return 1
       return 4
 
 
   def bar():
+      """
+      This is a multi-line
+      docstring
+      """
+
       if x:
           pass
           for l in x:
@@ -44,16 +51,25 @@ Execute(Blank lines should be added appropriately):
 
 Expect python(Newlines should be added):
   def foo():
+      """ This is a simple test docstring """
+
       return 1
 
 
   def bar():
+      '''This is another simple test docstring'''
+
       return 1
 
       return 4
 
 
   def bar():
+      """
+      This is a multi-line
+      docstring
+      """
+
       if x:
           pass
 
@@ -109,3 +125,43 @@ Expect python(extra newlines shouldn't be added to the main block):
 
   if __name__ == '__main__':
       main()
+
+
+Given python(A file with variables/docstring that start with a control statement):
+  def some():
+      """
+      This is a docstring that contains an
+      break control statement and also contains a
+      return something funny.
+      """
+
+      continue_some_var = True
+      forward_something = False
+
+      if (
+          continue_some_var and
+          forwarded_something
+      ):
+          return True
+
+
+Execute(Fix the file):
+  let g:ale_fixers = {'python': ['add_blank_lines_for_python_control_statements']}
+  ALEFix
+
+Expect python(Extra new lines are not added to the file):
+  def some():
+      """
+      This is a docstring that contains an
+      break control statement and also contains a
+      return something funny.
+      """
+
+      continue_some_var = True
+      forward_something = False
+
+      if (
+          continue_some_var and
+          forwarded_something
+      ):
+          return True