参数化 return 类型提示与 pandas 数据帧

Question

我有一个函数可以 return 不同的类型，例如 dict 或 pd.DataFrame。为了提供适当的类型提示，我天真地认为我可以这样做来实现这一点：

T = TypeVar("T")

def get_obj(key: str, typ: T = Any) -> T:
    # load obj using key
    return obj

但现在当我这样做时

df = get_obj(key="asd", typ=pd.DataFrame)

我的 linter 认为 df 是 class pd.DataFrame，而不是 pd.DataFrame 的实例。例如它说(variable) df Type[DataFrame]（而不是(variable) df DataFrame），为我提供了不同的实例化方法，当我做某事时很不高兴df.loc。

我知道我可以这样做：

@overload
def get_obj(key: str, is_df: Literal[False]) -> Any:
    ...

@overload
def get_obj(key: str, is_df: Literal[True]) -> pd.DataFrame:
    ...

def get_obj(key: str, is_df: bool = False):
    # load obj using key
    return obj

但我必须对所有可能的类型进行硬编码。

Answer 1

这是一个很酷的问题！我在下面有工作代码（它根据函数中给定的 typ 生成正确的智能感知）。

linter 没有得到它的原因是你给它一个通用类型，而不是 Type 类型。 typ 参数的内容毕竟是对 Type 的引用，而不是实例或字符串。

幸运的是，typing 库有 Type 类型，这有助于这个用例。

from typing import TypeVar, Type
import pandas as pd


item1: pd.DataFrame = pd.DataFrame()
item2: int = 2
item3: str = "HelloWorld!"

x = {
    "one": item1,
    "two": item2,
    "three": item3
}

T = TypeVar("T")

def get_obj(key: str, typ: Type[T]) -> T:
    # load obj using key
    obj: T = x[key]
    print(type(obj))
    return obj


y = get_obj(key="one", typ=pd.DataFrame)

注意，因为typ参数的内容是可变的，所以我们还是需要使用TypeVar。

以上代码将使 IDE 将变量 y 视为 pd.dataframe。如果你用 y = get_obj(key="one", typ=int) 赋值，你的 IDE 会认为对象是 int 类型，而实际上它是 pd.dataframe.

参数化 return 类型提示与 pandas 数据帧

Parametrized return type-hint with pandas DataFrames

python

dataframe

python-3.x

python-typing