Merge pull request #471 from birkenfeld/issue469

Make elpy-importmagic-fixup query the user for each unresolved symbol

elpy.el

@@ -2068,15 +2068,16 @@ prefix argument is given, prompt for a symbol from the user."
         (delete-region (point) (progn (forward-line (- end-line start-line)) (point)))
         (insert new-block)))))
 
-(defun elpy-importmagic--add-import-read-args ()
+(defun elpy-importmagic--add-import-read-args (&optional object prompt)
   (let* ((default-object (save-excursion
                            (let ((bounds (with-syntax-table python-dotty-syntax-table
                                            (bounds-of-thing-at-point 'symbol))))
                              (if bounds (buffer-substring (car bounds) (cdr bounds)) ""))))
-         (object-to-import (read-string "Object to import: " default-object))
+         (object-to-import (or object (read-string "Object to import: " default-object)))
          (possible-imports (elpy-rpc "get_import_symbols" (list buffer-file-name
                                                                 (elpy-rpc--buffer-contents)
-                                                                object-to-import))))
+                                                                object-to-import)))
+         (statement-prompt (or prompt "New import statement: ")))
     (cond
      ;; An elpy warning (i.e. index not ready) is returned as a string.
      ((stringp possible-imports)
@@ -2088,7 +2089,7 @@ prefix argument is given, prompt for a symbol from the user."
      ;; We have some candidates, let the user choose one.
      (t
       (let ((first-choice (car possible-imports))
-            (user-choice (completing-read "New import statement: " possible-imports)))
+            (user-choice (completing-read statement-prompt possible-imports)))
         (list (if (equal user-choice "") first-choice user-choice)))))))
 
 (defun elpy-importmagic-add-import (statement)
@@ -2100,10 +2101,21 @@ prefix argument is given, prompt for a symbol from the user."
       (elpy-importmagic--replace-block res))))
 
 (defun elpy-importmagic-fixup ()
+  "Query for new imports of unresolved symbols, and remove unreferenced imports.
+
+Also sort the imports in the import statement blocks."
   (interactive)
-  ;; get a new import statement block
-  (let* ((res (elpy-rpc "fixup_imports" (list buffer-file-name
-                                              (elpy-rpc--buffer-contents)))))
+  ;; get all unresolved names, and interactively add imports for them
+  (let* ((res (elpy-rpc "get_unresolved_symbols" (list buffer-file-name
+                                                       (elpy-rpc--buffer-contents)))))
+    (unless (stringp res)
+      (dolist (object res)
+        (let* ((prompt (format "How to import \"%s\": " object))
+               (choice (elpy-importmagic--add-import-read-args object prompt)))
+          (elpy-importmagic-add-import (car choice))))))
+  ;; now get a new import statement block (this also sorts)
+  (let* ((res (elpy-rpc "remove_unreferenced_imports" (list buffer-file-name
+                                                            (elpy-rpc--buffer-contents)))))
     (unless (stringp res)
       (elpy-importmagic--replace-block res))))
 

elpy/impmagic.py

@@ -59,9 +59,17 @@ def add_import(self, source, statement):
         start_line, end_line, import_block = imports.get_update()
         return start_line, end_line, import_block
 
-    def fixup_imports(self, source):
+    def get_unresolved_symbols(self, source):
         scope = importmagic.symbols.Scope.from_source(source)
-        unref, unres = scope.find_unresolved_and_unreferenced_symbols()
+        unres, unref = scope.find_unresolved_and_unreferenced_symbols()
+        return list(unres)
+
+    def remove_unreferenced_imports(self, source):
+        scope = importmagic.symbols.Scope.from_source(source)
+        unres, unref = scope.find_unresolved_and_unreferenced_symbols()
+        # Note: we do not supply "unres" to the call below, since we do
+        # not want to add imports without querying the user from which
+        # module symbols should be imported.
         start_line, end_line, import_block = importmagic.importer.get_update(
-            source, self.symbol_index, unref, unres)
+            source, self.symbol_index, set(), unref)
         return start_line, end_line, import_block

elpy/server.py

@@ -218,14 +218,21 @@ def rpc_add_import(self, filename, source, statement):
         source = get_source(source)
         return self.import_magic.add_import(source, statement)
 
-    def rpc_fixup_imports(self, filename, source):
-        """Automatically fixup imports in the file, by inserting import
-        statements.
+    def rpc_get_unresolved_symbols(self, filename, source):
+        """Return a list of unreferenced symbols in the module.
 
         """
         self._ensure_import_magic()
         source = get_source(source)
-        return self.import_magic.fixup_imports(source)
+        return self.import_magic.get_unresolved_symbols(source)
+
+    def rpc_remove_unreferenced_imports(self, filename, source):
+        """Remove unused import statements.
+
+        """
+        self._ensure_import_magic()
+        source = get_source(source)
+        return self.import_magic.remove_unreferenced_imports(source)
 
 
 def get_source(fileobj):

elpy/tests/test_impmagic.py

@@ -10,9 +10,11 @@
 
 TEST_SOURCE = '''# test file
 
-import sys
+import time
+import logging
 
 os.getcwd()
+time.sleep(1)
 '''
 
 
@@ -43,19 +45,26 @@ def test_add_import(self):
         start, end, newblock = self.importmagic.add_import(
             TEST_SOURCE, 'from mymod import AnUncommonName')
         self.assertEqual(start, 2)
-        self.assertEqual(end, 4)
+        self.assertEqual(end, 5)
         self.assertEqual(newblock.strip(),
-                         'import sys\n\nfrom mymod import AnUncommonName')
+                         'import logging\nimport time\n\nfrom mymod import AnUncommonName')
 
         start, end, newblock = self.importmagic.add_import(
             TEST_SOURCE, 'import mymod')
         self.assertEqual(start, 2)
-        self.assertEqual(end, 4)
-        self.assertEqual(newblock.strip(), 'import sys\n\nimport mymod')
+        self.assertEqual(end, 5)
+        self.assertEqual(newblock.strip(),
+                         'import logging\nimport time\n\nimport mymod')
+
+    def test_get_unresolved_symbols(self):
+        self.build_index()
+        symbols = self.importmagic.get_unresolved_symbols('x = a + b\ny = c.d')
+        self.assertEqual(sorted(symbols), ['a', 'b', 'c.d'])
 
-    def test_fixup(self):
+    def test_remove_unreferenced_imports(self):
         self.build_index()
-        start, end, newblock = self.importmagic.fixup_imports(TEST_SOURCE)
+        start, end, newblock = \
+            self.importmagic.remove_unreferenced_imports(TEST_SOURCE)
         self.assertEqual(start, 2)
-        self.assertEqual(end, 4)
-        self.assertEqual(newblock.strip(), 'import os')
+        self.assertEqual(end, 5)
+        self.assertEqual(newblock.strip(), 'import time')

elpy/tests/test_server.py

@@ -348,8 +348,10 @@ def test_should_call_importmagic(self):
             impmagic.get_import_symbols.assert_called_with("os")
             self.srv.rpc_add_import("filename", "source", "import os")
             impmagic.add_import.assert_called_with("source", "import os")
-            self.srv.rpc_fixup_imports("filename", "source")
-            impmagic.fixup_imports.assert_called_with("source")
+            self.srv.rpc_get_unresolved_symbols("filename", "source")
+            impmagic.get_unresolved_symbols.assert_called_with("source")
+            self.srv.rpc_remove_unreferenced_imports("filename", "source")
+            impmagic.remove_unreferenced_imports.assert_called_with("source")
 
 
 class TestGetSource(unittest.TestCase):

test/elpy-importmagic-fixup-test.el

@@ -0,0 +1,13 @@
+(ert-deftest elpy-importmagic-fixup-test ()
+  (elpy-testcase ()
+    (insert "# test
+import sys
+
+os.getcwd()
+")
+    (when (gethash "importmagic_version" (elpy-config--get-config))
+      (mletf* ((completing-read (prompt vals &optional func require-match default history)
+                                "import os"))
+        (elpy-importmagic-fixup)
+        (should (s-contains-p "import os\n" (buffer-string)))
+        (should-not (s-contains-p "import sys\n" (buffer-string)))))))