01-ibis-comp.py

import re

import tree_sitter_python as tspython

from IPython import get_ipython
from IPython.core.completer import SimpleCompletion, SimpleMatcherResult

from tree_sitter import Language, Parser


UNDERSCORE_PATTERN = re.compile(r".*\b_.$")

PY_LANGUAGE = Language(tspython.language())
parser = Parser(PY_LANGUAGE)
query = PY_LANGUAGE.query("(identifier) @name")


def ibis_underscore_completer(self, context) -> SimpleMatcherResult:
    if UNDERSCORE_PATTERN.match(context.token) is None:
        return self.python_matcher(context)

    import ibis.expr.types as ir

    ns = self.namespace
    tree = parser.parse(context.full_text.encode())
    root = tree.root_node
    captures = query.captures(root)
    name_nodes = captures["name"]
    attrs = next(
        (
            [f".{attr}" for attr in dir(obj) if not attr.startswith("_")]
            for node in reversed(name_nodes)
            if isinstance(obj := ns.get(node.text.decode()), ir.Table)
        ),
        [],
    )

    if not attrs:
        return self.python_matcher(context)

    _, fragment = self._attr_matches(context.line_with_cursor)

    return {
        "completions": [
            SimpleCompletion(text=attr, type="attribute") for attr in attrs
        ],
        "suppress": bool(attrs),
        "matched_fragment": fragment,
    }



ibis_underscore_completer.matcher_api_version = 2

ip = get_ipython()
ip.set_custom_completer(ibis_underscore_completer)

First you need to install tree sitter and tree sitter python bindings:

pip install tree-sitter tree-sitter-python

Stick this in ~/.ipython/profile_default/startup (or in the startup directory of the IPython profile you're using).