首页 > 解决方案 > 具有混合 dtypes 的 pandas DataFrame 的类型提示

问题描述

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

通常,如果我想提示一个函数的类型,它有一个 DataFrame 作为输入参数,我会这样做:

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

我似乎找不到的是如何键入提示具有混合 dtypes 的 DataFrame。DataFrame 构造函数仅支持完整 DataFrame 的类型定义。所以据我所知,dtypes 的变化只能在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 过程的机器学习应用程序,我认为类型提示很重要。

你用什么,有人在熊猫中输入提示他们的数据框吗?

标签: pythonpandastype-hinting

解决方案


我现在找到了看起来很有希望的 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 等静态类型检查器利用。有关更多详细信息,请参见此处的讨论。

但是,即使没有静态类型检查,我认为这是一个非常好的方向。


推荐阅读