Pandas DataFrame Accessor Type Hints

user_8039441 0
Viewed 491

Pandas allows you to extend its DataFrame class by using the pd.api.extensions.register_dataframe_accessor() decorator.

While this is functional, it doesn't offer any additional type hinting capabilities.

For example, I would expect the following to type check OK and even provide type hints

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

How can I get dataframe accessors to provide autocomplete?

pandas pycharm python-3.x types
1 Answers

This can be done somewhat hackishly using typing.TYPE_CHECKING and a bit of inheritance.

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

Unfortunately, PyCharm does not check the __annotations__ dictionary, or really do any dynamic type checking, so there doesn't appear to be any more universal solutions.

user_80394410
Related