Pandas DataFrame 访问器类型提示

Pandas 允许您使用以下方法扩展其 DataFrame 类pd.api.extensions.register_dataframe_accessor()装饰师。

虽然这是功能性的，但它不提供任何额外的类型提示功能。

例如，我希望以下内容能够键入 check OK，甚至提供类型提示

import pandas as pd
@pd.api.extensions.register_dataframe_accessor('dataset')
class Extension:
    def __init__(self, df: pd.DataFrame):
        self._df = df
    
    def foo(self, bar) -> str:
        return "foobar";

foo = pd.DataFrame({"foo":["bar"]})
foo.dataset.foo("bar")
    ^
    No Suggestions

如何让数据帧访问器提供自动完成功能？

这可以通过使用有点 hackish 来完成typing.TYPE_CHECKING和一点继承。

from typing import TYPE_CHECKING
import pandas as pd
@pd.api.extensions.register_dataframe_accessor('dataset')
class Extension:
    def __init__(self, df: pd.DataFrame):
        self._df = df
    
    def foo(self, bar) -> str:
        return "foobar";

if TYPE_CHECKING:
    class DataFrame(pd.DataFrame):
        dataset: Extension

foo: 'DataFrame' = pd.DataFrame({"foo":["bar"]})
#    ^ you have to do this every time you transform the DataFrame
foo.dataset.foo("bar")
#           ^ autocomplete is now provided
            

不幸的是，PyCharm 不检查__annotations__字典，或者真正进行任何动态类型检查，因此似乎没有任何更通用的解决方案。