使用pandas库的Python类型提示？

Question

让我们来看一个简单的函数，它接受一个字符串并返回一个数据帧：

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

如何使用Pythonic的方式为此函数添加类型提示？

如果我要求Python返回DataFrame的类型，它会返回 pandas.core.frame.DataFrame。但下面的代码并不能正常工作，因为它会告诉我pandas没有定义。

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

- Daniel

---

但是您正在使用pd别名，而且您可能可以定义自定义类型。 - Moses Koledoye

@MosesKoledoye 如果我尝试 pd.core.frame.DataFrame，我会得到一个 AttributeError 而不是 NameError。 - Daniel

我不是“Pythonicity”的权威，但我建议使用文档字符串（使用'''此函数接受一个inputType并返回一个outputType'''），这也是如果有人调用help(yourFunction)函数时将显示的内容。 - Chris

dataenforce 库可以用于检查数据框中的数据类型。https://github.com/CedricFR/dataenforce - 00schneider

在r/learnpython上相关内容：如何使用列指定pandas类型提示 - starball

Answer 1

为什么不直接使用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

- Georgy

---

它也不允许为特定的列指定数据类型，这可能非常有用。 - Philipp_Kats

目前还没有办法在类型提示中指定DataFrame列的数据类型，而且我还没有看到有关这方面的工作（如果我错了请纠正）。以下链接是与NumPy和数据类型相关的问题：Type hint for NumPy ndarray dtype?，你会发现那里也还没有实现此功能（https://github.com/numpy/numpy-stubs/issues/7）。 - Georgy

这在mypy中会出现错误：error: No library stub file for module 'pandas'。 - user2304916

@user2304916 请查看无法抑制“模块没有库存栏文件”的错误。 - Georgy

@Nesha25 这也类似于 list[int] 和 list[str] 之间的区别。如果没有类型参数告诉你列表中的内容，“里面”到底是什么，你就不知道可以合法地对其执行哪些操作。数据框架也是如此。另外，你举的“大于3的整数”的例子确实很不寻常，但是在类型系统中这种“值约束”并不罕见——例如考虑一个“非空指针”、“非零除数”或“带有已验证电子邮件地址的对象”。这些类型在许多地方都得到了应用。 - jonaslb

Answer 2

我目前正在进行以下工作：

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

我不知道这是否符合Pythonic风格，但作为类型提示，我觉得它足够易懂。

- Daniel

---

@Azat Ibrakov，您能详细解释一下您的评论吗？有时我不确定什么是“Pythonic”的概念。 - Tom Roth

我看到有人在给这个答案点踩。为了情境，这是我自己提出的问题并找到的解决方案，就所有意图和目的而言，它能够正常工作。上面更加符合Python风格的解决方案（我已接受为正确答案，但有其自身优点，请见评论）是在8个月之后才提供的。 - Daniel

这种写法不符合 Pythonic 的编程习惯，因为相比于此问题的标准答案而言，可读性更低且难以维护。由于类型路径在这里未经编译器验证，因此如果出现错误，程序将不会报错。这可能由于 TypeVar 参数的拼写错误或模块本身的更改导致。 - Alex

当我使用这个东西时，我收到了一个警告：TypeVar（）的参数必须是与其分配的变量名相等的字符串。 - Victor M Perez

@Azat Ibrakov，“Pythonic”和“非Pythonic”的争论对于许多“Pythonists”来说就像一个口头禅。我认为我们应该停止这种风格的争论。我从未听过Java开发人员使用这种类型的论证。在我看来，这个解决方案没有任何问题。 - uetoyo

这不是正确使用类型变量的方式。TypeVar 的存在是为了将两种类型联系起来（mypy文档）。您可能想要使用类型别名： PandasDataFrame = pandas.core.frame.DataFrame。 - decorator-factory

Answer 3

2023年11月更新

由于dataenforce的维护者停止了开发，我想更新一下这个答案。根据他的说法：

作为替代，我建议使用pandera，它提供了类似的功能（还有更多！）

---

原始答案

现在有一个pip包可以帮助解决这个问题。 https://github.com/CedricFR/dataenforce

您可以使用pip install dataenforce进行安装，并使用非常Pythonic的类型提示，例如：

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

- luksfarris

Answer 4

请看 pandera。

pandera提供了灵活而表达力强的API，用于对类似数据框的对象执行数据验证，以使数据处理流程更易读且更健壮。数据框包含pandera在运行时明确验证的信息。这在生产关键或可再现研究设置中非常有用。

---

pandera的优点是您还可以指定单个DataFrame列的数据类型。以下示例使用pandera来强制运行时实施一个包含整数列的DataFrame：

import pandas as pd
import pandera
from pandera.typing import DataFrame, Series

class Integers(pandera.SchemaModel):
    number: Series[int] 

@pandera.check_types
def my_fn(a: DataFrame[Integers]) -> None:
    pass

# This works
df = pd.DataFrame({"number": [ 2002, 2003]})
my_fn(df)

# Raises an exception
df = pd.DataFrame({"number": [ 2002.0, 2003]})
my_fn(df)

# Raises an exception
df = pd.DataFrame({"number": [ '2002', 2003]})
my_fn(df)

- Dvir Berebi

Answer 5

请查看这里给出的答案，该答案解释了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

- kevin_theinfinityfund

---

不幸的是，这被埋在底部。在2021年，这是最好的答案。还要注意链接答案后面的Daniel Malachov的评论（https://dev59.com/U1QJ5IYBdhLWcg3wIiS5#63446142）。 - user3897315

@user3897315 - 我不同意这是2021年最好的答案。如果您访问GitHub上的data-science-types存储库，您会发现该存储库已被归档，并且README（于2021年2月16日）已更新，加入以下说明："⚠️ 此项目大部分停止开发 ⚠️ Pandas团队和NumPy团队都正在将类型声明集成到他们的代码库中，我们不认为与他们竞争有意义。" - blthayer

我同意，但是在此之后，我没有看到pandas或numpy将在其路线图中推出这些功能的时间表或预计完成时间。 - kevin_theinfinityfund

Answer 6

这有点偏离原问题，但是可以在 @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='#')

- Keith