Find references gives a KeyError

[robmck1995]

Hi! Just switched our project over from pyright as the reference finding was massively superior.

However, just ran into this issue and not sure what the fix is.

I'm using this with pandas

Making the following reference request

{"id":1,"jsonrpc":"2.0","method":"textDocument/references","params":{"context":{"includeDeclaration":true},"position":{"character":8,"line":11731},"textDocument":{"uri":"file:///mnt/workspace/pandas/core/frame.py"}}}

Which corresponds to this function in pandas/core/frame.py

11730:    # error: Signature of "max" incompatible with supertype "NDFrame"
11731:    @overload  # type: ignore[override]
11732:    def max(
     : ________^
11733:        self,
11734:        *,
11735:        axis: Axis = ...,
11736:        skipna: bool = ...,
11737:        numeric_only: bool = ...,
11738:        **kwargs,
11739:    ) -> Series: ...

And I'm getting the following error

Error -32603: KeyError: <ForStmt: for i in range(1, nlev):             if lshape[i] == 0:                 stride = np.int64(0)             else:                 stride //= lshape[i]             out += labels[i] * stride@185,8>

Which appears to correspond to this section of code in pandas/core/sorting.py

185:        for i in range(1, nlev):
186:            if lshape[i] == 0:
187:                stride = np.int64(0)
188:            else:
189:                stride //= lshape[i]
190:            out += labels[i] * stride

I figure there may be some more details that would be helpful for debugging this which I'm happy to provide!

[davidhalter]

It would be very helpful if there was a minimal reproduction case (ideally without pandas). Because otherwise somebody else has to do that.

[robmck1995]

I'd be happy to find a reproduction case! Given the nature of the error (not clear to me given jedi novice and doesn't appear to show up much in the codebase) and the fact it appears to be a strange cross file issue, I thought maybe someone would have at least a guide on what might be going on.

I just doing really know where I would go on trying to reproduce this with a smaller example as I don't really know what is happening as the two code snippets don't seem very related.

Does that make sense?

[davidhalter]

Hmm it's a bit hard to tell without the traceback. With the traceback it would probably be quite easy. Can you access that somehow?

[robmck1995]

Sure thing! I apologize about the formatting, its the best I could get

{ jsonrpc: "2.0",
 id: 1, 
 method: None,
 params: None,
 result: None,
 error: {
   code: -32603,
   message: "KeyError: <ForStmt: for i in range(1, nlev):             if lshape[i] == 0:                 stride = np.int64(0)             else:                 stride //= lshape[i]             out += labels[i] * stride@185,8>",
   data: {
     "traceback": [
     String("  File "/app/venv/lib/python3.11/site-packages/pygls/protocol/json_rpc.py", line 266, in _handle_request
    self._execute_request(msg_id, handler, params)
"), String("  File "/app/venv/lib/python3.11/site-packages/pygls/protocol/json_rpc.py", line 188, in _execute_request
    self._send_response(msg_id, handler(params))
                                ^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi_language_server/server.py", line 428, in references
    names = jedi_script.get_references(*jedi_lines)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/api/helpers.py", line 487, in wrapper
    return func(self, line, column, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/api/__init__.py", line 429, in get_references
    return _references(**kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/api/__init__.py", line 423, in _references
    names = find_references(self._get_module_context(), tree_name, scope == 'file')
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/references.py", line 147, in find_references
    new = _dictionarize(_find_names(module_context, name_leaf))
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/references.py", line 74, in _find_names
    found_names = set(name.goto())
                      ^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/names.py", line 205, in goto
    values = infer_call_of_leaf(context, name, cut_own_trailer=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/helpers.py", line 104, in infer_call_of_leaf
    values = context.infer_node(base)
             ^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 224, in infer_node
    return infer_node(self, node)
           ^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 157, in infer_node
    return _infer_node_if_inferred(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 170, in _infer_node_if_inferred
    return _infer_node_cached(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/cache.py", line 44, in wrapper
    rv = function(obj, *args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 175, in _infer_node_cached
    return _infer_node(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 185, in _infer_node
    return infer_atom(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 305, in infer_atom
    return context.py__getattribute__(atom, position=position)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in py__getattribute__
    values = ValueSet.from_sets(name.infer() for name in names)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 430, in from_sets
    for set_ in sets:
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in <genexpr>
    values = ValueSet.from_sets(name.infer() for name in names)
                                ^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/names.py", line 281, in infer
    return tree_name_to_values(
           ^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/__init__.py", line 21, in wrapper
    return built_functions[public_name](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/stdlib.py", line 878, in wrapper
    return func(inference_state, context, tree_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/django.py", line 177, in wrapper
    result = func(inference_state, context, tree_name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 758, in tree_name_to_values
    types = infer_expr_stmt(context, node, tree_name)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 377, in infer_expr_stmt
    return _infer_expr_stmt(context, stmt, seek_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 410, in _infer_expr_stmt
    value_set = context.infer_node(rhs)
                ^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 224, in infer_node
    return infer_node(self, node)
           ^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 157, in infer_node
    return _infer_node_if_inferred(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 170, in _infer_node_if_inferred
    return _infer_node_cached(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/cache.py", line 44, in wrapper
    rv = function(obj, *args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 175, in _infer_node_cached
    return _infer_node(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 209, in _infer_node
    value_set = infer_trailer(context, value_set, trailer)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 279, in infer_trailer
    return atom_values.execute(args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 487, in execute
    return ValueSet.from_sets(c.inference_state.execute(c, arguments) for c in self._set)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 430, in from_sets
    for set_ in sets:
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 487, in <genexpr>
    return ValueSet.from_sets(c.inference_state.execute(c, arguments) for c in self._set)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/__init__.py", line 21, in wrapper
    return built_functions[public_name](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/stdlib.py", line 134, in wrapper
    return call()
           ^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/stdlib.py", line 107, in call
    return callback(value, arguments=arguments)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/pytest.py", line 29, in wrapper
    return callback(value, arguments)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/__init__.py", line 120, in execute
    value_set = value.py__call__(arguments=arguments)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/value/function.py", line 116, in py__call__
    return function_execution.infer()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/value/function.py", line 351, in infer
    return self.get_return_values()
           ^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/cache.py", line 44, in wrapper
    rv = function(obj, *args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/recursion.py", line 87, in wrapper
    result = func(self, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/value/function.py", line 242, in get_return_values
    value_set |= self.infer_node(children[1])
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 224, in infer_node
    return infer_node(self, node)
           ^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 157, in infer_node
    return _infer_node_if_inferred(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 170, in _infer_node_if_inferred
    return _infer_node_cached(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/cache.py", line 44, in wrapper
    rv = function(obj, *args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 175, in _infer_node_cached
    return _infer_node(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 185, in _infer_node
    return infer_atom(context, element)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 305, in infer_atom
    return context.py__getattribute__(atom, position=position)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in py__getattribute__
    values = ValueSet.from_sets(name.infer() for name in names)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 430, in from_sets
    for set_ in sets:
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in <genexpr>
    values = ValueSet.from_sets(name.infer() for name in names)
                                ^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/names.py", line 281, in infer
    return tree_name_to_values(
           ^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/__init__.py", line 21, in wrapper
    return built_functions[public_name](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/stdlib.py", line 878, in wrapper
    return func(inference_state, context, tree_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/django.py", line 177, in wrapper
    result = func(inference_state, context, tree_name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 758, in tree_name_to_values
    types = infer_expr_stmt(context, node, tree_name)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 377, in infer_expr_stmt
    return _infer_expr_stmt(context, stmt, seek_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 422, in _infer_expr_stmt
    left_values = context.py__getattribute__(name, position=stmt.start_pos)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in py__getattribute__
    values = ValueSet.from_sets(name.infer() for name in names)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/base_value.py", line 430, in from_sets
    for set_ in sets:
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 77, in <genexpr>
    values = ValueSet.from_sets(name.infer() for name in names)
                                ^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/names.py", line 281, in infer
    return tree_name_to_values(
           ^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/__init__.py", line 21, in wrapper
    return built_functions[public_name](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/stdlib.py", line 878, in wrapper
    return func(inference_state, context, tree_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/plugins/django.py", line 177, in wrapper
    result = func(inference_state, context, tree_name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 758, in tree_name_to_values
    types = infer_expr_stmt(context, node, tree_name)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 83, in wrapper
    return func(context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 377, in infer_expr_stmt
    return _infer_expr_stmt(context, stmt, seek_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/debug.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/syntax_tree.py", line 449, in _infer_expr_stmt
    with context.predefine_names(for_stmt, dct):
"), String("  File "/usr/lib/python3.11/contextlib.py", line 144, in __exit__
    next(self.gen)
"), String("  File "/app/venv/lib/python3.11/site-packages/jedi/inference/context.py", line 161, in predefine_names
    del predefined[flow_scope]
        ~~~~~~~~~~^^^^^^^^^^^^
")]}) }) }

[davidhalter]

Thanks. This is definitely a Jedi issue. Can you maybe try to reproduce it with a jedi.Script(path="....").get_references(line, column) call? You could then try to narrow it down to a small reproduction, because I'm pretty sure this is reproducible in 5-10 lines (within a single file).