Skip to content

Pandas 类型标注:高效且易维护代码的最佳实践

Updated on

Pandas 类型标注是 Python 开发者可以利用的一项强大工具,用来编写更高效、更易维护的代码。它主要是把 type hints(类型提示)和 pandas 这个常用的数据分析与处理库结合起来使用。本文的核心关键词是「pandas typing」,我们会介绍它在 Python 编程中的应用、带来的好处,以及如何借助它提升代码质量。

在 Python 中使用 type hints,尤其是配合 pandas,并不仅仅是一种潮流,而是正在 Python 社区中愈发普及的最佳实践。Type hints 相当于一种可执行的文档,能够帮助开发者弄清楚函数期望接收的数据类型以及返回的数据类型。这在处理 pandas 的 DataFrame 和 Series 时尤为有用,因为这些对象中包含的数据类型可能非常多样。

更新(pandas 2.1+):pandas 现在自带官方类型标注(py.typed),因此一般不再需要第三方 pandas-stubsmypypyright 等工具可以直接对 DataFrame / Series / Index API 做类型检查;建议固定到 pandas 2.2+ 以获得最完整的类型注解。

什么是 Pandas Typing?

Pandas typing 指的是将 type hints 和 pandas(一款 Python 中的数据分析与处理库)结合使用。它通过为函数参数和返回值标注预期的数据类型,帮助你在开发阶段更早地发现潜在错误。对于大型代码库,这种实践可以显著提升代码的可读性和可维护性。

Pandas typing 不只是给 pandas 代码随便加上一些类型标注,而是要充分利用 Python 的 typing 模块,让你的 pandas 代码更加健壮、更加容易理解。尤其在处理 pandas 的 DataFrame 和 Series 时,合理的类型标注可以帮助你保证使用的是正确的数据类型,避免运行时才暴露问题。

如何在 Pandas 中使用 typing 模块

Python 内置的 typing 模块专门用来书写类型提示。它为函数注解提供了一套标准的表示法,让你可以更明确地说明函数和方法的输入、输出类型。下面是一个将 typing 模块与 pandas 结合使用的示例:

from typing import List
import pandas as pd
 
def create_dataframe(data: List[dict]) -> pd.DataFrame:
    return pd.DataFrame(data)

在上面的例子中,create_dataframe 函数预期接收一个字典列表作为输入(data: List[dict]),并返回一个 pandas DataFrame(-> pd.DataFrame)。任何阅读代码的人都可以一眼看出这个函数期望的输入和返回值类型。

使用 Pandas 的 Pythonic 类型提示

针对 pandas 编写「Pythonic」的类型提示,就是用 Python 内置的 typing 模块对 pandas 相关代码做注解。这既能提升代码的可读性和可维护性,又可以在运行前发现潜在的类型问题。

例如,下面这个函数接收一个 pandas DataFrame,并返回一个 Series,你可以这样写类型提示:

from typing import List
import pandas as pd
 
def get_first_column(df: pd.DataFrame) -> pd.Series:
    return df.iloc[:, 0]

在这个例子里,get_first_column 函数预期接收一个 pandas DataFrame(df: pd.DataFrame),并返回一个 pandas Series(-> pd.Series)。这样任何人阅读代码时,都能清楚函数的输入输出类型。

在 pandas 中使用 Pythonic 的类型提示,可以帮助你确保自己在代码中使用了正确的数据类型,从而避免运行时错误,让代码更易理解和维护。

为指定列添加 Pandas 类型提示

在处理 pandas DataFrame 时,经常会遇到「不同列类型不同」的场景:某一列可能是整数,另一列是字符串,还有一列是布尔值。这种情况下,可以借助 pandas typing 清晰地描述 DataFrame 中各个列的类型。

示例如下:

from typing import Dict
import pandas as pd
 
def process_dataframe(df: pd.DataFrame) -> Dict[str, pd.Series]:
    # Process DataFrame and return a dictionary of Series
    result = {
        "int_column": df["int_column"].astype(int),
        "str_column": df["str_column"].astype(str),
        "bool_column": df["bool_column"].astype(bool),
    }
    return result

在这个例子中,process_dataframe 接收一个 pandas DataFrame,返回一个字典,其中每个键是列名,每个值是特定类型的 pandas Series。这让读者非常容易理解:每一列应该包含什么类型的数据。

使用 pandas typing 为特定列指定数据类型,有助于保证 DataFrame 中的列类型正确,从而减少运行时错误,也让后续对 DataFrame 的操作更加顺畅。

在 Pandas DataFrame 中使用类型提示的好处

在 pandas 代码中引入类型提示,可以显著提升代码质量。主要体现在:

  • 代码更易读:函数期望接收和返回的数据类型都是显式的;
  • 更易维护:在重构或协作开发时,类型信息是可靠的契约;
  • 更少运行时错误:许多类型相关的问题可以在静态检查阶段暴露出来。

例如,这里有个函数,它接收一个 pandas DataFrame,并返回一个只包含数值列的 DataFrame,可以这样写:

from typing import List
import pandas as pd
 
def select_numeric_columns(df: pd.DataFrame) -> pd.DataFrame:
    numeric_df = df.select_dtypes(include=[np.number])
    return numeric_df

在这个例子中,select_numeric_columns 函数预期接收一个 pandas DataFrame(df: pd.DataFrame),并返回一个只包含数值列的 DataFrame(-> pd.DataFrame)。看到签名的人就能立刻明白这个函数会做什么类型上的约束。

给已有的 Pandas 代码补充类型提示

如果你已经有一套使用 pandas 的代码,但里面没有任何类型提示,也完全可以逐步补充,以提升可读性和可维护性。大致步骤如下:

  1. typing 模块中导入需要的类型;
  2. 为每个函数或方法补充参数和返回值的类型注解;
  3. 使用 mypy 等类型检查工具跑一遍,验证类型标注是否正确。

下面是给已有 pandas 代码补充类型提示的一个简单示例:

# Before
import pandas as pd
 
def calculate_mean(df):
    return df.mean()
 
# After
from typing import List
import pandas as pd
 
def calculate_mean(df: pd.DataFrame) -> pd.Series:
    return df.mean()

如果你使用的是 pyright 这种较为严格的类型检查器,可以添加一个精简的 pyrightconfig.json 来收紧检查范围:

{
  "typeCheckingMode": "strict",
  "reportMissingTypeStubs": "warning"
}

在 pandas 2.2+ 的环境下,上面的配置无需额外的 stub 包就可以正常工作。

在添加类型注解后的版本中,calculate_mean 函数期望接收一个 pandas DataFrame,并返回一个 pandas Series,这一点变得非常清晰。


常见问题解答

pandas DataFrame 和 pandas Series 有什么区别?

pandas DataFrame 是一个带标签的二维数据结构,每一列可以是不同的数据类型。它类似于电子表格、SQL 表,或者由若干 Series 构成的字典。pandas Series 则是一个带标签的一维数组,可以容纳任意数据类型。

pandas API 中的新式类型标注是什么?

pandas API 中的新式类型标注,是指通过 Python 内置的 typing 模块,为 pandas 相关代码添加类型注解。这有助于提升代码的可读性和可维护性,也能在运行前捕获潜在的类型错误。

如何提升 pandas 类型标注相关的效率?

想要在使用 pandas 类型标注时兼顾性能与效率,可以从以下几方面入手:

  • 合理为函数参数和返回值添加 type hints;
  • 使用 mypy 等类型检查工具定期验证类型注解;
  • 在 DataFrame 中为特定列指定明确的数据类型,减少隐式转换和不必要的类型混用。