有关熊猫的 Python 类型提示？

让我们来看一个简单的函数，它接受一个 str 并返回一个数据帧:

import pandas as pd
def csv_to_df(path):
return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

向这个函数添加类型提示的推荐 Python 方法是什么？

如果我向 python 询问 DataFrame 的类型，它会返回 pandas.core.frame.DataFrame。 下面的内容不会起作用，因为它会告诉我，熊猫是没有定义的。

 def csv_to_df(path: str) -> pandas.core.frame.DataFrame:
return pd.read_csv(path, skiprows=1, sep='\t', comment='#')
84351 次浏览
小开

我目前正在做以下工作:

from typing import TypeVar
PandasDataFrame = TypeVar('pandas.core.frame.DataFrame')
def csv_to_df(path: str) -> PandasDataFrame:
return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

结果是:

> help(csv_to_df)
Help on function csv_to_df in module __main__:


csv_to_df(path:str) -> ~pandas.core.frame.DataFrame

我不知道这句话有多玄乎，但我发现，作为一个类型提示，它是可以理解的。

小开
最佳答案

为什么不用 pd.DataFrame呢？

import pandas as pd
def csv_to_df(path: str) -> pd.DataFrame:
return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

结果是一样的:

> help(csv_to_df)
Help on function csv_to_df in module __main__:
csv_to_df(path:str) -> pandas.core.frame.DataFrame
小开

这偏离了最初的问题，但是基于@dangom 的回答，使用 TypeVar和@Georgy 的评论，没有办法在类型提示中为 DataFrame 列指定数据类型，你可以使用这样一个简单的解决方案在 DataFrame 中指定数据类型:

from typing import TypeVar
DataFrameStr = TypeVar("pandas.core.frame.DataFrame(str)")
def csv_to_df(path: str) -> DataFrameStr:
return pd.read_csv(path, skiprows=1, sep='\t', comment='#')
小开

现在有一个 pip 软件包可以帮助解决这个问题。 Https://github.com/cedricfr/dataenforce

您可以使用 pip install dataenforce安装它，并使用非常 Python 化的类型提示，如:

def preprocess(dataset: Dataset["id", "name", "location"]) -> Dataset["location", "count"]:
pass
小开

查看给定 给你的答案，该答案解释了包 data-science-types的用法。

pip install data-science-types

演示

# program.py


import pandas as pd


df: pd.DataFrame = pd.DataFrame({'col1': [1,2,3], 'col2': [4,5,6]}) # OK
df1: pd.DataFrame = pd.Series([1,2,3]) # error: Incompatible types in assignment

用同样的方法运行 mypy:

$ mypy program.py


提交答案