如何在VS Code中让pandas注册的DataFrame访问器显示文档与类型提示

我通过register_dataframe_accessor装饰器扩展了Pandas DataFrame的API，代码能正常运行，属性和方法都可以调用，但在VS Code里看不到自定义访问器的文档字符串和类型提示；而在Jupyter Notebook/JupyterLab中显示正常，测试虚拟方法时也存在同样差异。请问如何让VS Code支持自定义DataFrame访问器的文档与类型提示？

相关代码如下：

import pandas as pd
import numpy as np

@pd.api.extensions.register_dataframe_accessor("geo")
class GeoAccessor:
    def __init__(self, pandas_obj):
        self._validate(pandas_obj)
        self._obj = pandas_obj

    @staticmethod
    def _validate(obj):
        # 验证是否包含latitude和longitude列
        if "latitude" not in obj.columns or "longitude" not in obj.columns:
            raise AttributeError("必须包含'latitude'和'longitude'列。")

    @property
    def center(self):
        # 返回该DataFrame的地理中心点
        lat = self._obj.latitude
        lon = self._obj.longitude
        return (float(lon.mean()), float(lat.mean()))

    def plot(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        """绘制数据到地图上，例如使用Cartopy
        """
        return dataframe

ds = pd.DataFrame(
    {"longitude": np.linspace(0, 10), "latitude": np.linspace(0, 20)}
)
ds.geo.plot()

解决方法 #

1. 添加静态类型注解扩展DataFrame类型 #

VS Code的语言服务器（如Pylance）是静态分析代码，无法识别Pandas运行时注册的访问器。可以通过TYPE_CHECKING标记创建一个子类，给类型检查器明确的提示：

import pandas as pd
import numpy as np
from typing import TYPE_CHECKING

# 仅在类型检查阶段生效
if TYPE_CHECKING:
    # 定义包含geo访问器的DataFrame子类
    class GeoDataFrame(pd.DataFrame):
        @property
        def geo(self) -> 'GeoAccessor': ...

@pd.api.extensions.register_dataframe_accessor("geo")
class GeoAccessor:
    def __init__(self, pandas_obj):
        self._validate(pandas_obj)
        self._obj = pandas_obj

    @staticmethod
    def _validate(obj):
        if "latitude" not in obj.columns or "longitude" not in obj.columns:
            raise AttributeError("必须包含'latitude'和'longitude'列。")

    @property
    def center(self):
        """返回DataFrame的地理中心点坐标"""
        lat = self._obj.latitude
        lon = self._obj.longitude
        return (float(lon.mean()), float(lat.mean()))

    def plot(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        """在地图上绘制数据（例如使用Cartopy）
        参数:
            dataframe: 待绘制的DataFrame对象
        返回:
            处理后的DataFrame
        """
        return dataframe

# 给变量添加类型标注，让VS Code识别geo访问器
ds: GeoDataFrame = pd.DataFrame(
    {"longitude": np.linspace(0, 10), "latitude": np.linspace(0, 20)}
)
ds.geo.plot()

2. 配置Pylance增强类型检查 #

确保VS Code使用Pylance作为默认语言服务器，然后在设置中开启相关选项：

• 打开设置（快捷键Ctrl+,），搜索python.analysis.typeCheckingMode，设置为basic或strict
• 开启python.analysis.inlayHints.enable以显示更详细的类型提示

3. 全局扩展DataFrame的__getattr__类型提示 #

如果有多个自定义访问器，可通过给pd.DataFrame添加__getattr__的类型注解来统一处理：

from typing import Any, TypeVar
import pandas as pd

DataFrameT = TypeVar("DataFrameT", bound=pd.DataFrame)

if TYPE_CHECKING:
    def __getattr__(self: DataFrameT, name: str) -> Any: ...
    pd.DataFrame.__getattr__ = __getattr__

这样类型检查器会允许访问任意自定义属性，并根据实际类型推断提示。

内容的提问来源于stack exchange，提问作者Alix TC