类型提示 Pandas DataFrame 内容和列

Question

我正在编写一个返回 PandasDataFrame对象的函数。我希望有某种类型来暗示它DataFrame包含哪些列，而不仅仅是文档中的规范，因为我觉得这将使最终用户更容易读取数据。

在编辑 Python 文件和编辑 Jupyter Notebook 时，是否有办法输入DataFrameVisual Studio Code 和 PyCharm 等不同工具支持的提示内容？

一个示例函数：

def generate_data(bunch, of, inputs) -> pd.DataFrame:
      """Massages the input to a nice and easy DataFrame.

      :return:
           DataFrame with columns a(int), b(float), c(string), d(us dollars as float)
      """

Answer 1

截至目前（2023 年 4 月），pandas 强类型最强大的项目是pandera。不幸的是，它提供的功能非常有限，远远达不到我们想要的。DataFrame

\n

pandera以下是如何在您的情况下使用^\xe2\x80\xa0的示例：

\n

import pandas as pd\nimport pandera as pa\nfrom pandera.typing import DataFrame\n\nclass MySchema(pa.DataFrameModel):\n    a: int\n    b: float\n    c: str = pa.Field(nullable=True)  # For example, allow None values\n    d: float    # US dollars\n\nclass OtherSchema(pa.DataFrameModel):\n    year: int = pa.Field(ge=1900, le=2050)\n\n\ndef generate_data() -> DataFrame[MySchema]:\n    df = pd.DataFrame({\n        "a": [1, 2, 3],\n        "b": [10.0, 20.0, 30.0],\n        "c": ["A", "B", "C"],\n        "d": [0.1, 0.2, 0.3],\n    })\n\n    # Runtime verification here, throws on schema mismatch\n    strongly_typed_df = DataFrame[MySchema](df)\n    return strongly_typed_df\n\ndef transform(input: DataFrame[MySchema]) -> DataFrame[OtherSchema]:\n    # This demonstrates that you can use strongly\n    # typed column names from the schema\n    df = input.filter(items=[MySchema.a]).rename(\n            columns={MySchema.a: OtherSchema.year}\n    )\n\n    return DataFrame[OtherSchema](df) # This will throw on range validation!\n\n\ndf1 = generate_data()\ndf2 = transform(df1)\ntransform(df2)   # mypy prints error here - incompatible type!\n

\n

mypy您可以在最后一行看到产生静态类型检查错误：

\n

\n

优点和局限性的讨论

\n

使用 pandera 我们得到 \xe2\x80\x93

\n

\n
1. 清晰易读（dataclass风格）DataFrame的模式定义以及将它们用作类型提示的能力。
\n
2. 运行时模式验证。模式可以定义比类型更多的约束（请参阅year下面的示例和pandera文档以了解更多信息）。
\n
3. mypy对静态类型检查的实验支持。
\n

\n

我们仍然怀念的\xe2\x80\x93

\n

\n
1. 用于列级验证的完整静态类型检查。
\n
2. 任何支持列名称自动完成的 IDE。
\n
3. 模式声明的内联语法，我们必须在使用每个模式之前将其显式定义为单独的类。
\n

\n

更多示例

\n

Pandera 文档 - https://pandera.readthedocs.io/en/stable/dataframe_models.html

\n

类似的问题-带有混合数据类型的 pandas DataFrame 的类型提示

\n

其他打字项目

\n

pandas-stubs是一个活跃的项目，为 pandas 公共 API 提供类型声明，它比 pandas 本身包含的类型存根更丰富。但它不为列级模式提供任何设施。

\n

有很多过时的库与此和 pandas 一般类型相关 - dataenforce、data-science-types、python-type-stubs

\n

_{\xe2\x80\xa0pandera提供了两种不同的 API，它们似乎同样强大 - 基于对象的 API 和基于类的 API。我在这里演示了后者。}

\n