diff --git a/pyls/plugins/importmagic_lint.py b/pyls/plugins/importmagic_lint.py
index 15035f3a..b8cc7f00 100644
--- a/pyls/plugins/importmagic_lint.py
+++ b/pyls/plugins/importmagic_lint.py
@@ -129,7 +129,7 @@ def pyls_code_actions(config, document, context):
                 # These tend to be terrible
                 continue
 
-            actions.append(generate_add_action(document.source, index, module, variable))
+            actions.append(generate_add_action(document, index, module, variable))
 
     return actions