From ccbf6cd7b32db8fd338fb6767bdb73bdfd10b4cf Mon Sep 17 00:00:00 2001
From: Stephen Macke <stephen.macke@gmail.com>
Date: Tue, 3 Oct 2023 14:29:00 -0700
Subject: [PATCH] Fix numpy go-to-definition by taking it off autoimport list
 for this case (#447)

---
 pylsp/plugins/definition.py      | 22 ++++++++++++++++------
 test/plugins/test_definitions.py | 12 ++++++++++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/pylsp/plugins/definition.py b/pylsp/plugins/definition.py
index ffc1b00b..53eda915 100644
--- a/pylsp/plugins/definition.py
+++ b/pylsp/plugins/definition.py
@@ -3,6 +3,9 @@
 from __future__ import annotations
 import logging
 from typing import Any, Dict, List, TYPE_CHECKING
+
+import jedi
+
 from pylsp import hookimpl, uris, _utils
 
 if TYPE_CHECKING:
@@ -43,12 +46,19 @@ def pylsp_definitions(
     settings = config.plugin_settings("jedi_definition")
     code_position = _utils.position_to_jedi_linecolumn(document, position)
     script = document.jedi_script(use_document_path=True)
-    definitions = script.goto(
-        follow_imports=settings.get("follow_imports", True),
-        follow_builtin_imports=settings.get("follow_builtin_imports", True),
-        **code_position,
-    )
-    definitions = [_resolve_definition(d, script, settings) for d in definitions]
+    auto_import_modules = jedi.settings.auto_import_modules
+
+    try:
+        jedi.settings.auto_import_modules = []
+        definitions = script.goto(
+            follow_imports=settings.get("follow_imports", True),
+            follow_builtin_imports=settings.get("follow_builtin_imports", True),
+            **code_position,
+        )
+        definitions = [_resolve_definition(d, script, settings) for d in definitions]
+    finally:
+        jedi.settings.auto_import_modules = auto_import_modules
+
     follow_builtin_defns = settings.get("follow_builtin_definitions", True)
     return [
         {
diff --git a/test/plugins/test_definitions.py b/test/plugins/test_definitions.py
index f0e9ffef..c366e8ca 100644
--- a/test/plugins/test_definitions.py
+++ b/test/plugins/test_definitions.py
@@ -36,6 +36,9 @@ def my_func():
 inception = my_list[2]
 
 inception()
+
+import numpy
+numpy.ones
 """
 
 
@@ -89,6 +92,15 @@ def test_definition_with_multihop_inference_goto(config, workspace):
     )
 
 
+def test_numpy_definition(config, workspace):
+    # Over numpy.ones
+    cursor_pos = {"line": 29, "character": 8}
+
+    doc = Document(DOC_URI, workspace, DOC)
+    defns = pylsp_definitions(config, doc, cursor_pos)
+    assert len(defns) > 0, defns
+
+
 def test_builtin_definition(config, workspace):
     # Over 'i' in dict
     cursor_pos = {"line": 8, "character": 24}