Fix python import regex

Author: AndrewRadev

autoload/sj/python.vim

@@ -92,7 +92,7 @@ function! sj#python#JoinTuple()
 endfunction
 
 function! sj#python#SplitImport()
-  let import_pattern = '^from \%(.*\) import \zs.*$'
+  let import_pattern = '^\s*from \%(.*\) import \zs.*$'
 
   normal! 0
   if search(import_pattern, 'Wc', line('.')) <= 0
@@ -141,7 +141,7 @@ function! sj#python#SplitImport()
 endfunction
 
 function! sj#python#JoinImportWithNewlineEscape()
-  let import_pattern = '^from \%(.*\) import .*\\\s*$'
+  let import_pattern = '^\s*from \%(.*\) import .*\\\s*$'
   if getline('.') !~ import_pattern
     return 0
   endif
@@ -160,7 +160,7 @@ function! sj#python#JoinImportWithNewlineEscape()
 endfunction
 
 function! sj#python#JoinImportWithRoundBrackets()
-  let import_pattern = '^from \%(.*\) import \zs('
+  let import_pattern = '^\s*from \%(.*\) import \zs('
   if search(import_pattern, 'Wc') <= 0
     return 0
   endif

spec/plugin/python_spec.rb

@@ -62,37 +62,50 @@
   end
 
   specify "imports" do
-    set_file_contents 'from foo import bar, baz'
+    set_file_contents <<~EOF
+      def surrounding_function():
+          from foo import bar, baz
+    EOF
 
+    vim.search('from foo')
     split
 
     assert_file_contents <<~EOF
-      from foo import bar,\\
-              baz
+      def surrounding_function():
+          from foo import bar,\\
+                  baz
     EOF
 
     join
-    assert_file_contents 'from foo import bar, baz'
+    assert_file_contents <<~EOF
+      def surrounding_function():
+          from foo import bar, baz
+    EOF
 
     vim.command 'let b:splitjoin_python_import_style = "round_brackets"'
     split
 
     assert_file_contents <<~EOF
-      from foo import (bar,
-                       baz)
+      def surrounding_function():
+          from foo import (bar,
+                           baz)
     EOF
 
     join
-    assert_file_contents 'from foo import bar, baz'
+    assert_file_contents <<~EOF
+      def surrounding_function():
+          from foo import bar, baz
+    EOF
 
     vim.command 'let b:splitjoin_python_brackets_on_separate_lines = 1'
     split
 
     assert_file_contents <<~EOF
-      from foo import (
-          bar,
-          baz
-      )
+      def surrounding_function():
+          from foo import (
+              bar,
+              baz
+          )
     EOF
 
     join