jedi: add update_user_modules() function

This function will be used to fix import completion of user module
names in Pybricks Code.

Issue: pybricks/support#759

Author: dlech

jedi/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ## Unreleased
 
+### Added
+- Added `update_user_modules()` function for filtering on user modules.
+
 ### Fixed
 - Fixed code completion for builtin types.
 - Fixed code completion for names starting with `_`.

jedi/src/pybricks_jedi/__init__.py

@@ -2,6 +2,7 @@
 import json
 import re
 from enum import IntEnum
+from typing import Iterable
 
 import docstring_parser
 import jedi
@@ -176,6 +177,8 @@
 
 MICROPY_NOT_SUPPORTED_DUNDER = {"__doc__", "__package__"}
 
+user_modules = set()
+
 # Types from monaco editor
 
 
@@ -326,7 +329,7 @@ def _is_pybricks(c: Completion) -> bool:
 
     # filter out packages/modules that are not included in Pybricks firmware
     if c.type == "module" or c.type == "namespace":
-        return c.full_name in PYBRICKS_CODE_PACKAGES
+        return c.full_name in PYBRICKS_CODE_PACKAGES or c.full_name in user_modules
 
     # filter subset of builtins
     if c.module_name == "builtins" and c.type != "keyword":
@@ -541,3 +544,15 @@ def get_signatures(code: str, line: int, column: int) -> str:
     """
     signatures = jedi.Script(code).get_signatures(line, column - 1)
     return json.dumps(_map_signatures(signatures))
+
+
+def update_user_modules(names: Iterable[str]) -> None:
+    """
+    Updates the set of user module names used for filtering.
+
+    Args:
+        names:
+            An iterable of module names.
+    """
+    user_modules.clear()
+    user_modules.update(names)

jedi/tests/test_complete_import.py

@@ -6,7 +6,9 @@
 """
 
 import json
-from pybricks_jedi import CompletionItem, complete
+
+import pytest
+from pybricks_jedi import CompletionItem, complete, update_user_modules
 
 
 def test_from():
@@ -27,6 +29,32 @@ def test_from():
     ]
 
 
+@pytest.fixture
+def user_modules():
+    update_user_modules(["jedi", "pytest"])
+    yield
+    update_user_modules([])
+
+
+def test_from_with_user_modules(user_modules):
+    code = "from "
+    completions: list[CompletionItem] = json.loads(complete(code, 1, len(code) + 1))
+    assert [c["insertText"] for c in completions] == [
+        "jedi",
+        "micropython",
+        "pybricks",
+        "pytest",
+        "uerrno",
+        "uio",
+        "ujson",
+        "umath",
+        "urandom",
+        "uselect",
+        "ustruct",
+        "usys",
+    ]
+
+
 def test_from_pybricks_import():
     code = "from pybricks import "
     completions: list[CompletionItem] = json.loads(complete(code, 1, len(code) + 1))