Don't get confused by multiline strings that look like imports

Wilfred · Wilfred · commit d0d77c518705 · 2017-08-08T11:39:33.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@
   unsaved or not visiting a file.
 * `pyimport-insert-missing` now steps over shebangs when choosing
   where to insert new imports.
+* Fixed an issue with `pyimport-insert-missing` using the contents of
+  multiline strings that looked like imports.
 
 ## v1.0
 
diff --git a/pyimport.el b/pyimport.el
@@ -46,9 +46,22 @@
   "Return non-nil if the current line is the last in the buffer."
   (looking-at (rx (0+ not-newline) buffer-end)))
 
+(defun pyimport--in-string-p ()
+  "Return non-nil if point is inside a string."
+  (nth 3 (syntax-ppss)))
+
 (defun pyimport--buffer-lines (buffer)
-  (with-current-buffer buffer
-    (s-split "\n" (buffer-string))))
+  "Return all the lines in BUFFER, ignoring lines that are within a string."
+  (let (lines)
+    (with-current-buffer buffer
+      (save-excursion
+        (goto-char (point-min))
+        (unless (pyimport--in-string-p)
+          (push (pyimport--current-line) lines))
+        (while (zerop (forward-line 1))
+          (unless (pyimport--in-string-p)
+            (push (pyimport--current-line) lines)))))
+    (nreverse lines)))
 
 (defun pyimport--import-lines (buffer)
   "Return all the lines in this Python BUFFER that look like imports."
diff --git a/test/pyimport-test.el b/test/pyimport-test.el
@@ -185,3 +185,12 @@
    (equal
     (pyimport--extract-simple-import "import foo as bar" "bar")
     "import foo as bar")))
+
+(ert-deftest pyimport-buffer-lines ()
+  "We should ignore multiline strings, as they may not contain valid imports."
+  (let (lines)
+    (with-temp-buffer
+      (insert "x = \"\"\"\nfrom foo\n\"\"\"\n")
+      (setq lines (pyimport--buffer-lines (current-buffer))))
+    (should (not
+             (-contains-p lines "from foo")))))
