Ввод фрейма данных Polars с проверкой Pandera и Mypy ⇐ Python

[Anonymous]

Я рассматриваю pandera для реализации строгой типизации в моем проекте с использованием полярных фреймов данных.
Я не понимаю, как правильно типизировать свои функции.< /p>
В качестве примера:

Код: Выделить всё

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 определяется как реализация DataFrameBase в коде pandera.
Как я могу исправить свой код, чтобы возвращаемый тип LazyFrame[MyModel] работает?

Подробнее здесь: https://stackoverflow.com/questions/790 ... validation