具有混合数据类型的 pandas DataFrame 的类型提示

Question

我一直在寻找 pandas DataFrame 的健壮类型提示，但似乎找不到任何有用的东西。这个问题仅仅触及表面Pandas 的 Pythonic 类型提示？

通常，如果我想暗示以 DataFrame 作为输入参数的函数的类型，我会这样做：

import pandas as pd 
def func(arg: pd.DataFrame) -> int: 
     return 1

我似乎找不到的是如何使用mixed dtypes输入提示DataFrame。 DataFrame 构造函数仅支持完整 DataFrame 的类型定义。因此，据我所知，数据类型的更改只能在使用

pd.DataFrame().astype(dtypes={})

函数之后发生。

这在这里可行，但对我来说似乎不太Pythonic

import datetime
def func(arg: pd.DataFrame(columns=['integer', 'date']).astype(dtype={'integer': int, 'date': datetime.date})) -> int:
    return 1

我遇到了这个包：https://pypi.org/project/dataenforce/，其中包含这样的示例：

def process_data(data: Dataset["id": int, "name": object, "latitude": float, "longitude": float])
  pass

这看起来很有希望，但遗憾的是该项目已经过时且存在缺陷。

作为一名数据科学家，在构建具有长 ETL 过程的机器学习应用程序时，我认为类型提示非常重要。

你使用什么？有人在 pandas 中输入暗示他们的数据框吗？

Answer 1

我现在找到了看起来非常有前途的pandera库：

https://github.com/pandera-dev/pandera

它允许用户创建模式并使用这些模式来创建详细检查。来自他们的文档：

https://pandera.readthedocs.io/en/stable/schema_models.html

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


class InputSchema(pa.SchemaModel):
    year: Series[int] = pa.Field(gt=2000, coerce=True)
    month: Series[int] = pa.Field(ge=1, le=12, coerce=True)
    day: Series[int] = pa.Field(ge=0, le=365, coerce=True)

class OutputSchema(InputSchema):
    revenue: Series[float]

@pa.check_types
def transform(df: DataFrame[InputSchema]) -> DataFrame[OutputSchema]:
    return df.assign(revenue=100.0)


df = pd.DataFrame({
    "year": ["2001", "2002", "2003"],
    "month": ["3", "6", "12"],
    "day": ["200", "156", "365"],
})

transform(df)

invalid_df = pd.DataFrame({
    "year": ["2001", "2002", "1999"],
    "month": ["3", "6", "12"],
    "day": ["200", "156", "365"],
})
transform(invalid_df)

还有他们的注释：

由于 pandas 库的当前限制（请参阅此处的讨论），pandera 注释仅用于运行时验证，不能被 mypy 等静态类型检查器利用。有关更多详细信息，请参阅此处的讨论。

但是，即使没有静态类型检查，我认为这正在朝着一个非常好的方向发展。

Answer 2

虽然不是 Pandas，但 StaticFrame（基于不可变数据模型构建的替代 DataFrame 库，我是其中的作者）允许 DataFrame 的完整类型提示，包括混合柱状类型。如果需要，您可以轻松地在 StaticFrame 和 Pandas 之间进行转换。

例如，要键入提示上面的

func

，我们可以将

arg

定义为具有任意类型索引、字符串列标签（以

sf.Frame

形式给出）和两列数据值的 StaticFrame

sf.Index[np.str_]

指定为

np.int64

和

np.object_

数组（最后一个包含

datetime.date

对象）。

from typing import Any
import numpy as np
import static_frame as sf

def func(arg: sf.Frame[Any, sf.Index[np.str_], np.int64, np.object_]) -> int:
    return 1

更多示例请参见此处。