Improve how jedi handles numpy

Gerard Vermeulen · Gerard Vermeulen · commit 08c14ea53996 · 2022-10-31T18:27:19.000+01:00
diff --git a/CONFIGURATION.md b/CONFIGURATION.md
@@ -16,6 +16,7 @@ This server can be configured using the `workspace/didChangeConfiguration` metho
 | `pylsp.plugins.flake8.indentSize` | `integer` | Set indentation spaces. | `null` |
 | `pylsp.plugins.flake8.perFileIgnores` | `array` of `string` items | A pairing of filenames and violation codes that defines which violations to ignore in a particular file, for example: `["file_path.py:W305,W304"]`). | `[]` |
 | `pylsp.plugins.flake8.select` | `array` of unique `string` items | List of errors and warnings to enable. | `null` |
+| `pylsp.plugins.jedi.auto_import_modules` | `array` of `string` items | List of module names for jedi.settings.auto_import_modules. | `["numpy"]` |
 | `pylsp.plugins.jedi.extra_paths` | `array` of `string` items | Define extra paths for jedi.Script. | `[]` |
 | `pylsp.plugins.jedi.env_vars` | `object` | Define environment variables for jedi.Script and Jedi.names. | `null` |
 | `pylsp.plugins.jedi.environment` | `string` | Define environment for jedi.Script and Jedi.names. | `null` |
diff --git a/pylsp/config/schema.json b/pylsp/config/schema.json
@@ -87,6 +87,14 @@
       "uniqueItems": true,
       "description": "List of errors and warnings to enable."
     },
+    "pylsp.plugins.jedi.auto_import_modules": {
+      "type": "array",
+      "default": ["numpy"],
+      "items": {
+        "type": "string"
+      },
+      "description": "List of module names for jedi.settings.auto_import_modules."
+    },
     "pylsp.plugins.jedi.extra_paths": {
       "type": "array",
       "default": [],
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
@@ -14,6 +14,8 @@
 
 log = logging.getLogger(__name__)
 
+DEFAULT_AUTO_IMPORT_MODULES = ["numpy"]
+
 # TODO: this is not the best e.g. we capture numbers
 RE_START_WORD = re.compile('[A-Za-z_0-9]*$')
 RE_END_WORD = re.compile('^[A-Za-z_0-9]*')
@@ -252,6 +254,8 @@ def jedi_script(self, position=None, use_document_path=False):
 
         if self._config:
             jedi_settings = self._config.plugin_settings('jedi', document_path=self.path)
+            jedi.settings.auto_import_modules = jedi_settings.get('auto_import_modules',
+                                                                  DEFAULT_AUTO_IMPORT_MODULES)
             environment_path = jedi_settings.get('environment')
             extra_paths = jedi_settings.get('extra_paths') or []
             env_vars = jedi_settings.get('env_vars')
