如何让自定义类在VSCode-Jupyter环境中支持键自动补全？

我希望自定义的mydict类能像原生dict或pd.DataFrame一样支持键自动补全，但目前实现的类无法触发键的补全：

class mydict:

    def __init__(self, _dict) -> None:
        self._dict = _dict

    def __getitem__(self, key):
        return self._dict[key]

原生dict中输入键“aa”时能触发自动补全，但这个自定义类做不到，我没找到相关解决方案。

解决方案 #

要让自定义类支持键自动补全，核心是让IDE能识别到类包含的键列表，可通过以下方式实现：

1. 重写__dir__方法 #

IDE的自动补全逻辑通常会读取对象__dir__方法返回的内容，我们可以把字典的键加入到返回列表中：

class mydict:
    def __init__(self, _dict) -> None:
        self._dict = _dict

    def __getitem__(self, key):
        return self._dict[key]

    def __dir__(self):
        # 返回类本身的属性方法 + 字典的键
        return super().__dir__() + list(self._dict.keys())

2. 支持点语法访问（可选） #

如果希望像访问属性一样用obj.aa代替obj['aa']，可以实现__getattr__方法，让补全体验更接近原生对象：

class mydict:
    def __init__(self, _dict) -> None:
        self._dict = _dict

    def __getitem__(self, key):
        return self._dict[key]

    def __getattr__(self, name):
        try:
            return self._dict[name]
        except KeyError:
            raise AttributeError(f"'mydict' object has no attribute '{name}'") from None

    def __dir__(self):
        return super().__dir__() + list(self._dict.keys())

3. 增强类型提示（可选） #

如果需要更精准的类型推断，可结合typing模块的泛型和TypedDict，让IDE提前识别键的类型：

from typing import Generic, TypeVar, TypedDict

T = TypeVar('T', bound=TypedDict)

class mydict(Generic[T]):
    def __init__(self, _dict: T) -> None:
        self._dict = _dict

    def __getitem__(self, key):
        return self._dict[key]

    def __getattr__(self, name):
        try:
            return self._dict[name]
        except KeyError:
            raise AttributeError(f"'mydict' object has no attribute '{name}'") from None

    def __dir__(self):
        return super().__dir__() + list(self._dict.keys())

使用时定义具体的TypedDict类型，IDE就能精准提示对应键：

class MyDictType(TypedDict):
    aa: int
    bb: str

obj = mydict(MyDictType(aa=1, bb="test"))
# 输入obj. 会自动提示aa和bb

内容的提问来源于stack exchange，提问作者Stupid Boy