auto-import completions for a submodule of a package when the package is already imported

[KevinMusgrave]

I'm familiar with sending requests like textDocument/signatureHelp and textDocument/completion.

If I have basedpyright.analysis.autoImportCompletions set to true, what request type should I send to get import completions?

Thanks!

[DetachHead]

textDocument/completion should work. is it not working for you?

which editor are you using? or are you using your own LSP client?

[KevinMusgrave]

I don't think it's working for me. I'm using my own LSP client. I'll try to put together a simple reproducible script.

[KevinMusgrave]

Here's the script:

import asyncio
import json
import uuid
import platform


class MinimalLanguageServer:
    def __init__(self):
        self.server = None
        self.initialized = False
        self.completion_request_id = None
        
    async def start_server(self):
        print("Starting server...")
        self.server = await asyncio.create_subprocess_shell(
            "basedpyright-langserver --stdio",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        
    async def send_message(self, message):
        request = json.dumps(message).encode("utf-8")
        content_length = len(request)
        header = f"Content-Length: {content_length}\r\n\r\n".encode("utf-8")
        self.server.stdin.write(header + request)
        await self.server.stdin.drain()

    async def listen_for_responses(self):
        while True:
            header = await self.server.stdout.readline()
            if not header:
                break
                
            if header.startswith(b"Content-Length: "):
                content_length = int(header.decode("utf-8").split(": ")[1])
                await self.server.stdout.readline()  # Empty line
                content = await self.server.stdout.read(content_length)
                response = json.loads(content)
                
                # Only print completion responses
                if "id" in response and response["id"] == self.completion_request_id:
                    print(json.dumps(response, indent=2))
                
                # Handle initialization silently
                if not self.initialized and "result" in response and "capabilities" in response["result"]:
                    self.initialized = True
                    await self.send_message({
                        "jsonrpc": "2.0",
                        "method": "initialized",
                        "params": {}
                    })

async def main():
    ls = MinimalLanguageServer()
    await ls.start_server()
    
    # Initialize request
    init_params = {
        "capabilities": {},
        "workspaceFolders": [{"uri": "file:///workspace"}],
    }
    await ls.send_message({
        "jsonrpc": "2.0",
        "id": str(uuid.uuid4()),
        "method": "initialize",
        "params": init_params
    })
    
    # Open document
    test_code = "import sklearn\nsklearn.datasets"
    doc_uri = "untitled:/test.py"
    if platform.system() == "Windows":
        doc_uri += "/"  # needed on Windows
    await ls.send_message({
        "jsonrpc": "2.0",
        "method": "textDocument/didOpen",
        "params": {
            "textDocument": {
                "uri": doc_uri,
                "languageId": "python",
                "version": 1,
                "text": test_code
            }
        }
    })
    
    # Send completion request
    ls.completion_request_id = str(uuid.uuid4())
    await ls.send_message({
        "jsonrpc": "2.0",
        "id": ls.completion_request_id,
        "method": "textDocument/completion",
        "params": {
            "textDocument": {"uri": doc_uri},
            "position": {"line": 1, "character": 16}
        }
    })
    
    # Listen for responses
    await ls.listen_for_responses()

if __name__ == "__main__":
    asyncio.run(main())

The response I get is:

{
  "jsonrpc": "2.0",
  "id": "4b4591d6-393f-488e-9b8c-fe255b922c3b",
  "result": {
    "items": [],
    "isIncomplete": true
  }
}

I know it has access to the correct python environment, because if I request completion results for:

import sklearn.datasets\nsklearn.datasets.m

then I get correct completion results like "make_blobs" and "make_biclusters" which are part of the sklearn.datasets module.

For the import completion, I'm expecting:

I am able to get import completions in some other situations, though I'm not sure how relevant they are. For example, if the text is just "os", I get lots of results like:

{
        "label": "ONE_SIXTH",
        "kind": 21,
        "data": {
          "uri": "untitled:/test.py",
          "position": {
            "line": 0,
            "character": 2
          },
          "funcParensDisabled": true,
          "autoImportText": "from colorsys import ONE_SIXTH",
          "symbolLabel": "ONE_SIXTH"
        },
        "sortText": "12.9999.ONE_SIXTH.08.colorsys",
        "detail": "Auto-import",
        "labelDetails": {
          "description": "colorsys"
        },
        "documentation": {
          "kind": "plaintext",
          "value": "from colorsys import ONE_SIXTH"
        },
        "textEdit": {
          "range": {
            "start": {
              "line": 0,
              "character": 0
            },
            "end": {
              "line": 0,
              "character": 2
            }
          },
          "newText": "ONE_SIXTH"
        },

[DetachHead]

there are two issues here.

the first is that requesting completions at the end of the completed module name won't return anything because the name is already there:

it works if you add a .:

but those completions suck because there's another problem here. in your script the code you're requesting completions for is:

import sklearn
sklearn.datasets

the above code imports the sklearn package instead of sklearn.datasets. normally, importing sklearn wouldn't give you access to sklearn.datasets because that module isn't imported inside sklearn/__init__.py. the only reason it works at runtime is because sklearn/__init__.py defines a top-level __getattr__ function that imports all of its submodules dynamically. the return type of __getattr__ is used by pyright when accessing any unknown attribute of sklearn:

import sklearn
reveal_type(sklearn.datasets) # ModuleType | Any

therefore the type information is lost which is why you aren't seeing any useful completions:

if you import the datasets module directly, the module's type doesn't get lost and you get your completions:

import sklearn.datasets
reveal_type(sklearn.datasets) # Module("sklearn.datasets")

i fixed these two issues in your script and now it returns the correct completions:

    # Open document
-   test_code = "import sklearn\nsklearn.datasets"
+   test_code = "import sklearn.datasets\nsklearn.datasets."
    doc_uri = "untitled:/test.py"
    if platform.system() == "Windows":
        doc_uri += "/"  # needed on Windows
    await ls.send_message({
        "jsonrpc": "2.0",
        "method": "textDocument/didOpen",
        "params": {
            "textDocument": {
                "uri": doc_uri,
                "languageId": "python",
                "version": 1,
                "text": test_code
            }
        }
    })
    
    # Send completion request
    ls.completion_request_id = str(uuid.uuid4())
    await ls.send_message({
        "jsonrpc": "2.0",
        "id": ls.completion_request_id,
        "method": "textDocument/completion",
        "params": {
            "textDocument": {"uri": doc_uri},
-           "position": {"line": 1, "character": 16}
+           "position": {"line": 1, "character": 17}
        }
    })

[KevinMusgrave]

I understand that you have to do import sklearn.datasets in order to get useful module/attribute completions for sklearn.datasets.

But I'd like the import completion for sklearn.datasets. That is, given the following code:

import sklearn

sklearn.datasets

it should suggest importing sklearn.datasets, which is the behavior I get in VSCode using Pylance.

[DetachHead]

ok so it looks like pylance supports completions for a module inside a package if the package's __init__.py is imported:

but pyright doesn't: