使用 pandera 和 mypy 验证输入极坐标数据框

Question

我正在考虑

pandera

使用

polars

数据帧来实现我的项目的强类型。

我对如何正确输入函数感到困惑。

举个例子：

import polars as pl
import pandera.polars as pa
from pandera.typing.polars import LazyFrame as PALazyFrame

class MyModel(pa.DataFrameModel):
    a: int
    class Config:
        strict = True


def foo(
    f: pl.LazyFrame    
) -> PALazyFrame[MyModel]:
    # Our input is unclean, probably coming from pl.scan_parquet on some files
    # The validation is dummy here
    return MyModel.validate(f.select('a'))

如果我调用

mypy

，它将返回以下错误

error: Incompatible return value type (got "DataFrameBase[MyModel]", expected "LazyFrame[MyModel]") 

当然，我可以修改我的签名以指定返回类型

DataFrameBase[MyModel]

，但我会失去返回 LazyFrame 的精度。

进一步，LazyFrame 被定义为在

pandera

代码中实现 DataFrameBase。

如何修复我的代码以使返回类型 LazyFrame[MyModel] 正常工作？

Answer 1

当底层库可能无法很好地表达类型时，这通常是一个问题 - 幸运的是，有几种方法可以解决这个问题：

1.演员之路

正如评论中所讨论的，使用

typing.cast

始终是一种选择。如果外部库不能生成足够具体的类型，这通常是我的选择 - 它比使用

type:ignore

好得多，并且允许您在其他类型良好的代码库中“弥合差距”。例如

import polars as pl
import pandera.polars as pa
from pandera.typing.polars import LazyFrame as PALazyFrame
import typing
class MyModel(pa.DataFrameModel):
    a: int
    class Config:
        strict = True

def foo(
    f: pl.LazyFrame    
) -> PALazyFrame[MyModel]:
    # Our input is unclean, probably coming from pl.scan_parquet on some files
    # The validation is dummy here
    return typing.cast(PALazyFrame[MyModel],MyModel.validate(f.select('a')))

正如所提到的，有时必须手动调整转换类型 - 并且这必须在可能的多个位置进行验证。

2.方法论

假设我们需要在代码中的很多地方使用这个模型，您可能希望将“转换”推得离最终用户更远一些。强制转换并没有真正消失，但它允许我们将其放在可以高度重用的地方，并减少代码库中强制转换的数量（始终是一个好目标！）。请注意，原始库中的底层代码确实使用了此方法的强制转换，因此我们实际上只是将其重新强制转换为略有不同的内容。

在下面的示例中，有一种专门用于验证惰性帧的新方法 - 它的操作方式与常规验证相同，只是它采用 LazyFrame 并输出 PALazyFrame：

import polars as pl
import pandera.polars as pa
from pandera.typing.polars import LazyFrame as PALazyFrame
from typing import Optional, Self, cast

class MyDataFrameModel(pa.DataFrameModel):
    @classmethod
    def validate_lazy(
        cls,
        check_obj: pl.LazyFrame,
        head: Optional[int]=None, 
        tail: Optional[int]=None, 
        sample: Optional[int]=None, 
        random_state: Optional[int]=None, 
        lazy: bool=False, 
        inplace: bool=False
    ) -> PALazyFrame[Self]:
        return cast(PALazyFrame[Self], cls.validate(
            check_obj,
            head,
            tail,
            sample,
            random_state,
            lazy,
            inplace
        ))

class MyModel(MyDataFrameModel):
    a: int
    class Config:
        strict = True

def foo(
    f: pl.LazyFrame    
) -> PALazyFrame[MyModel]:
    # Our input is unclean, probably coming from pl.scan_parquet on some files
    # The validation is dummy here
    return MyModel.validate_lazy(f.select('a'))

我最初考虑简单地用另一种更通用且允许此用例的方法覆盖原始验证方法，但我发现：

a) 很难表达正确的输出类型

b) 无论如何，现代 Python 都禁止以不兼容的方式覆盖继承模型中的方法。

3.公关之道

如果做不到这一点，您唯一的选择就是请求更改底层库。可能有一种方法可以以更通用的方式表达此方法，以适应您的用例，但是请记住，如果没有“更高种类的类型”的存在，有些类型结构根本无法表达，这目前 Python 中不存在。我怀疑这可能是这样的用例之一。

希望这有帮助！