Detect multiline imports and don't show snippets for them

pyls/plugins/jedi_completion.py

@@ -39,6 +39,9 @@
     'statement': lsp.CompletionItemKind.Keyword,
 }
 
+# Types of parso nodes for which snippet is not included in the completion
+_IMPORTS = ('import_name', 'import_from')
+
 
 @hookimpl
 def pyls_completions(config, document, position):
@@ -51,8 +54,8 @@ def pyls_completions(config, document, position):
 
     settings = config.plugin_settings('jedi_completion', document_path=document.path)
     should_include_params = settings.get('include_params')
-    include_params = (use_snippets(document, position) and
-                      snippet_support and should_include_params)
+    include_params = (snippet_support and should_include_params and
+                      use_snippets(document, position))
     return [_format_completion(d, include_params) for d in definitions] or None
 
 
@@ -63,13 +66,19 @@ def use_snippets(document, position):
     This returns `False` if a completion is being requested on an import
     statement, `True` otherwise.
     """
-    lines = document.source.split('\n')
-    act_line = lines[position['line']]
-    tokens = parso.parse(act_line)
-    act_statement = tokens.children[0].get_code()
-    if act_statement.startswith('import') or act_statement.startswith('from'):
-        return False
-    return True
+    line = position['line']
+    lines = document.source.split('\n', line)
+    act_lines = [lines[line][:position['character']]]
+    line -= 1
+    while line > -1:
+        act_line = lines[line]
+        if act_line.rstrip().endswith('\\'):
+            act_lines.insert(0, act_line)
+            line -= 1
+        else:
+            break
+    tokens = parso.parse('\n'.join(act_lines).split(';')[-1].strip())
+    return tokens.children[0].type not in _IMPORTS
 
 
 def _format_completion(d, include_params=True):

test/plugins/test_completion.py

@@ -166,3 +166,37 @@ def test_snippets_completion(config):
     completions = pyls_jedi_completions(config, doc, com_position)
     out = 'defaultdict(${1:default_factory}, ${2:iterable}, ${3:kwargs})$0'
     assert completions[0]['insertText'] == out
+
+
+def test_multiline_snippets(config):
+    document = 'from datetime import\\\n date,\\\n datetime \na=date'
+    doc = Document(DOC_URI, document)
+    config.capabilities['textDocument'] = {
+        'completion': {'completionItem': {'snippetSupport': True}}}
+    config.update({'plugins': {'jedi_completion': {'include_params': True}}})
+
+    position = {'line': 1, 'character': 5}
+    completions = pyls_jedi_completions(config, doc, position)
+    assert completions[0]['insertText'] == 'date'
+
+    position = {'line': 2, 'character': 9}
+    completions = pyls_jedi_completions(config, doc, position)
+    assert completions[0]['insertText'] == 'datetime'
+
+
+def test_multistatement_snippet(config):
+    config.capabilities['textDocument'] = {
+        'completion': {'completionItem': {'snippetSupport': True}}}
+    config.update({'plugins': {'jedi_completion': {'include_params': True}}})
+
+    document = 'a = 1; from datetime import date'
+    doc = Document(DOC_URI, document)
+    position = {'line': 0, 'character': len(document)}
+    completions = pyls_jedi_completions(config, doc, position)
+    assert completions[0]['insertText'] == 'date'
+
+    document = 'from datetime import date; a = date'
+    doc = Document(DOC_URI, document)
+    position = {'line': 0, 'character': len(document)}
+    completions = pyls_jedi_completions(config, doc, position)
+    assert completions[0]['insertText'] == 'date(${1:year}, ${2:month}, ${3:day})$0'