Papermill 扩展开发实战：通过 Entry Points 自定义 I/O 处理器与执行引擎

Papermill 在开箱即用地支持本地文件、S3、GCS、HDFS 等多种读写来源和默认的 nbclient 本地执行引擎之外，还提供了一套基于 Python Entry Points 的插件机制，允许开发者在不修改 papermill 源码的前提下，为其接入任意存储后端（如 SFTP 服务器）或自定义执行逻辑（如远程执行、逐单元格计时统计）。本文将以仓库 docs/extending-overview.rst、docs/extending-entry-points.rst 为核心脉络，结合 papermill/iorw.py、papermill/engines.py 等源码实现，完整讲解 I/O Handler 与 Engine 的接口约定、entry point 注册方式，并给出两个可复制运行的完整实战示例，使读者能够独立为 papermill 开发并交付自己的插件。

扩展机制总览：Papermill 的插件化设计

用 papermill 运行一个 notebook 时，背后其实只发生四件事（见 docs/extending-overview.rst）：

1. 读取 notebook 文件；
2. 将文件内容转换为 notebook 的 Python 对象（nbformat 的 NotebookNode）；
3. 执行该 notebook；
4. 将执行后的 notebook 写回文件。

其中步骤 1、3、4 正是扩展点：通过 entry points，你可以编写自己的工具来接管"读取"（I/O Handler）、"执行"（Engine）与"写回"（I/O Handler）。步骤 2 由 papermill 内部完成，通常无需干预。

所谓 entry points，是 Python 打包规范中"已安装的分发包向外界通告自身组件，供其他代码发现和使用"的机制——典型如 console_scripts 生成命令行包装器，以及 Pygments 通过它加载第三方语法高亮插件。Papermill 在运行时正是通过 entrypoints 库扫描两类 entry point 组：

• papermill.io：实现输入/输出（I/O）处理的 handler；
• papermill.engine：实现执行逻辑的 engine。

从源码可以印证这一点：papermill/iorw.py 中的 PapermillIO.register_entry_points() 调用 entrypoints.get_group_all("papermill.io") 并逐个注册；papermill/engines.py 中的 PapermillEngines.register_entry_points() 则调用 entrypoints.get_group_all("papermill.engine")。两个注册方法都在模块导入时被自动执行（iorw.py 与 engines.py），因此第三方插件只要以正常方式安装到环境中，papermill 启动时即会自动发现。

如果觉得仅靠新增 handler 和 engine 仍不够，而是想改进项目本身的一些根本性设计，则可以参与 papermill 的贡献开发，详见下文"参与 papermill 核心开发"一节（对应 docs/extending-developing.rst）。

开发新的 I/O Handler

I/O Handler 的接口约定：四个必须实现的方法

Papermill 中"读取输入 notebook"由 I/O Handler 管理，正是它们让 papermill 不仅能访问本地文件系统，还能访问 S3 等远程服务；同样，"将执行后的 notebook 写回"也由 I/O Handler 负责。因此，I/O Handler 是 papermill 与任意存储后端对接的统一抽象层。

编写自己的 I/O Handler，就是编写一个实现了以下四个类方法的类（CustomIO 仅为示意名）：

方法	签名	作用
read	CustomIO.read(file_path)	返回文件内容（字符串）
write	CustomIO.write(file_content, file_path)	写入文件，无返回值
pretty_path	CustomIO.pretty_path(path)	返回美化后的路径，用于日志与展示
listdir	CustomIO.listdir(path)	返回路径列表，用于目录浏览

原文档特别提醒：如果你的 handler 只用于写（例如只面向发布平台），不打算支持 read 等操作，那么应当实现该方法，并在被调用时抛出异常（如 NotImplementedError），而不是省略它——这样接口保持完整，行为上又明确表达了"此能力不受支持"。

这一约定与仓库内置 handler 的实现完全一致。例如 papermill/iorw.py 的 LocalHandler 实现了全部四个方法；而只读性质更强的 handler 则主动对不支持的方法抛异常，如 GithubHandler.write 抛 PapermillException('write is not supported by GithubHandler')（iorw.py）、StreamHandler.listdir 抛 PapermillException('listdir is not supported by Stream Handler')（iorw.py）。这些实现可以作为编写"部分能力受限" handler 的现成参考。

确保 handler 被 papermill 发现：pyproject.toml 注册

开发完 handler 类之后，需要在插件的 pyproject.toml 中声明 papermill 的 entry point。方法是在文件中加入 [project.entry-points."papermill.io"] 段：

[project.entry-points."papermill.io"]
"sftp://" = "papermill_sftp:SFTPHandler"

这一行的语义是：当传入的文件路径以 sftp:// 开头时，papermill 就使用 papermill_sftp 包中导入的 SFTPHandler 类来处理该路径的读写。等号左边是路径前缀，等号右边是类名及其导入来源（格式为 包名:类名）。

从源码看，这套匹配逻辑实现在 PapermillIO.get_handler 中（papermill/iorw.py）：它按注册顺序遍历 self._handlers，一旦 path.startswith(scheme) 命中即返回对应 handler；若全部未命中，则回退到 local handler（即本地文件系统），连本地 handler 都没有时才抛出 PapermillException。另外要注意 register 是 LIFO 顺序插入（iorw.py），后注册的 scheme 会优先匹配，设计同名前缀覆盖时需留意这一点。

传统上，papermill I/O handler 的 entry point 名采用 URL 前缀形式。仓库内置注册的 handler（注意：这些是 papermill 内部注册，并非通过 entry point 加载，但接口与约定相同）包括（见 papermill/iorw.py）：

• local → LocalHandler：本地文件系统；
• s3:// → S3Handler：Amazon S3；
• adl:// → ADLHandler：Azure Data Lake；
• abs:// → ABSHandler：Azure Blob Storage；
• http://、https:// → HttpHandler：HTTP/HTTPS 读写；
• gs:// → GCSHandler：Google Cloud Storage（写入时带限流重试）；
• hdfs:// → HDFSHandler：Hadoop 文件系统；
• http://github.com/、https://github.com/ → GithubHandler：GitHub 内容读取；
• - → StreamHandler：标准输入/输出流。

可以推断，任何以 URL 风格前缀命名的新 handler 都能与这套体系自然共存。

实战示例：完整的 SFTP I/O Handler

下面按原文档的演示，从零构建一个可读写 SFTP 服务器的 handler，目标是支持这样的命令行用法：

papermill sftp://my_ftp_server.co.uk/input.ipynb sftp://my_ftp_server.co.uk/output.ipynb

项目结构如下：

papermill_sftp
    |- pyproject.toml
    |- src
        |- papermill_sftp
            |- __init__.py

在 src/papermill_sftp/__init__.py 中实现 handler：读取时先把远端文件下载到临时目录再读入；写入时先写到临时文件再上传；pretty_path 直接原样返回路径；listdir 暂不实现，按接口约定抛异常。原文档示例代码中未显式导入 pathlib、tempfile、urllib.parse，且 cnopts 需按你的 pysftp 环境配置，这里补齐 import 并给出说明，使代码可直接运行：

import os
import pathlib
import tempfile
import urllib.parse

import pysftp

sftp_username = os.getenv('SFTP_USERNAME')
sftp_password = os.getenv('SFTP_PASSWORD')
# 根据你的环境配置主机密钥校验选项，例如：
# cnopts = pysftp.CnOpts()
# cnopts.hostkeys = None   # 仅测试环境使用，生产环境请校验 host key
cnopts = pysftp.CnOpts()


class SFTPHandler:

    @classmethod
    def read(cls, path):
        """
        Read a notebook from an SFTP server.
        """
        parsed_url = urllib.parse.urlparse(path)
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_file = pathlib.Path(tmpdir) / pathlib.Path(parsed_url.path).name
            with pysftp.Connection(
                parsed_url.hostname,
                username=sftp_username,
                password=sftp_password,
                port=(parsed_url.port or 22),
                cnopts=cnopts,
            ) as sftp:
                sftp.get(parsed_url.path, str(tmp_file))
            return tmp_file.read_text()

    @classmethod
    def write(cls, file_content, path):
        """
        Write a notebook to an SFTP server.
        """
        parsed_url = urllib.parse.urlparse(path)
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_file = pathlib.Path(tmpdir) / "output.ipynb"
            tmp_file.write_text(file_content)
            with pysftp.Connection(
                parsed_url.hostname,
                username=sftp_username,
                password=sftp_password,
                port=(parsed_url.port or 22),
                cnopts=cnopts,
            ) as sftp:
                sftp.put(str(tmp_file), parsed_url.path)

    @classmethod
    def pretty_path(cls, path):
        return path

    @classmethod
    def listdir(cls, path):
        raise NotImplementedError

配套的 pyproject.toml 完整内容如下（注意 [project.entry-points."papermill.io"] 段，以及 setuptools 的 src 布局声明）：

[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "papermill_sftp"
version = "0.1"
description = "An SFTP I/O handler for papermill."
authors = [
    {name = "My Name", email = "[e-mail]"}
]
dependencies = ["pysftp"]

[project.urls]
Repository = "https://github.com/my_username/papermill_sftp.git"

[project.entry-points."papermill.io"]
"sftp://" = "papermill_sftp:SFTPHandler"

[tool.setuptools]
packages = ["papermill_sftp"]
package-dir = {"" = "src"}

安装该插件后，papermill 执行时会检查输入、输出路径是否以 sftp:// 开头，若命中则调用 papermill_sftp 中的 SFTPHandler 完成读与写。整个交互链路为：sftp:// 路径 → PapermillIO.get_handler 前缀匹配（iorw.py）→ SFTPHandler.read/write → 继续走 papermill 的读取/写回流程。

开发新的执行引擎（Engine）

Engine 基类与 execute_managed_notebook 接口

Papermill 的 engine 是"能够执行一个 notebook 的 Python 对象"。默认实现 NBClientEngine 接收一个 notebook 对象并在本机执行（其背后是 nbclient 的 PapermillNotebookClient，见 papermill/engines.py）。通过编写自定义 engine，你可以把执行交给远程服务器，或在执行后对 notebook 做后处理（例如注入额外的输出单元格）。

自定义 engine 需要继承 papermill.engines.Engine 基类，并实现类方法 execute_managed_notebook，其调用签名要与父类保持一致：

class CustomEngine(papermill.engines.Engine):

    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):
        pass

这里需要澄清一个容易误解的细节：原文档将 nb_man 描述为 nbformat.NotebookNode 对象，但从源码看（papermill/engines.py），Engine.execute_notebook 会先把 notebook 包装进 NotebookExecutionManager（其内部 nb 属性才是 NotebookNode），再把该管理器实例传给 execute_managed_notebook。NotebookExecutionManager 封装了统一的执行状态管理：notebook_start 初始化并清空 papermill 元数据、cell_start/cell_exception/cell_complete 逐单元格更新状态与耗时、notebook_complete 收尾并强制保存，还内置了进度条与自动保存（autosave_cell_every，默认 30 秒，见 engines.py）。因此自定义 engine 只需聚焦"如何逐单元格执行"，而元数据记录、保存等横切逻辑可全部复用。

基类 Engine 的默认 execute_managed_notebook 直接抛出 NotImplementedError（engines.py），强制子类实现。

实战示例：记录每个代码单元格耗时的 Timing Engine

原文档提供了一个完整的演示：实现一个自定义 engine，把每个代码单元格的执行耗时作为额外输出注入到该单元格的 outputs 开头。因为 papermill 本身已在单元格元数据中记录 start_time / end_time，我们直接复用默认引擎 NBClientEngine，再借助 nbformat 的 new_output 构造输出节点。

项目结构：

papermill_timing
    |- pyproject.toml
    |- src
        |- papermill_timing
            |- __init__.py

src/papermill_timing/__init__.py 内容如下：

from datetime import datetime
from papermill.engines import NBClientEngine
from nbformat.v4 import new_output

class CustomEngine(NBClientEngine):

    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):

        # call the papermill execution engine:
        super().execute_managed_notebook(nb_man, kernel_name, **kwargs)

        for cell in nb_man.nb.cells:

            if cell.cell_type == "code" and cell.execution_count is not None:
                start = datetime.fromisoformat(cell.metadata.papermill.start_time)
                end = datetime.fromisoformat(cell.metadata.papermill.end_time)
                output_message = f"Execution took {(end - start).total_seconds():.3f} seconds"
                output_node = new_output("display_data", data={"text/plain": [output_message]})
                cell.outputs = [output_node] + cell.outputs

实现要点：先调用 super().execute_managed_notebook(...) 走完默认的本地执行流程，此时 cell.metadata.papermill.start_time / end_time 已被 NotebookExecutionManager.cell_start / cell_complete 写入（这两个回调分别见 engines.py 与 engines.py）；随后遍历代码单元格，计算耗时差，并把新构造的 display_data 输出节点插到原有 outputs 之前。

确保 engine 被 papermill 发现：pyproject.toml 注册

自定义 engine 需要以 papermill.engine 为前缀注册为 entry point，引用我们刚实现的类。以 papermill_timing 包中的 CustomEngine 为例：

[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "papermill_timing"
version = "0.1"
description = "A papermill engine that logs additional timing information about code."
authors = [
    {name = "My Name", email = "[e-mail]"}
]
dependencies = ["papermill", "nbformat"]

[project.urls]
Repository = "https://github.com/my_username/papermill_timing.git"

[project.entry-points."papermill.engine"]
timer_engine = "papermill_timing:CustomEngine"

[tool.setuptools]
packages = ["papermill_timing"]
package-dir = {"" = "src"}

注册后，用户即可通过命令行参数 --engine timer_engine 选用该引擎（该参数在 papermill/cli.py 中定义为"执行 notebook 时使用的引擎名称"）：

papermill input.ipynb output.ipynb --engine timer_engine

PapermillEngines.get_engine 会按名称在已注册的引擎字典中查找，找不到时抛出 PapermillException(f"No engine named '{name}' found")（engines.py）。仓库内置注册了 None 与 'nbclient' 两个名称，均指向 NBClientEngine（engines.py），自定义名称与内置名称互不冲突。

下图左侧为使用自定义 timing engine 执行后的 notebook，右侧为使用标准 engine 执行的结果：每个代码单元格顶部都多出了我们注入的耗时输出（配图源自 docs/extending-entry-points.rst）：

关于 nb_man 回调的更轻量写法

如果你的自定义 engine 不想走完整的 nbclient 执行流程，而只是想在默认流程之外做点事情，可以像测试用例 CellCallbackEngine 那样直接操纵回调：在 papermill/tests/test_engines.py 中，一个继承 Engine 的最小实现只调用 nb_man.cell_start(cell) 与 nb_man.cell_complete(cell)，配合 Engine.execute_notebook 的包装即可完成元数据的完整记录与保存。这验证了：引擎实现者的核心职责只是决定"什么时候执行、执行什么"，状态机与持久化由 NotebookExecutionManager 统一兜底。

参与 papermill 核心开发

当扩展需求超出"新增 I/O handler 与执行 handler"的范畴，涉及项目根本性改进时，可以选择直接向 papermill 贡献代码（详见 docs/extending-developing.rst）。仓库内已提供 CONTRIBUTING.md、CODE_OF_CONDUCT.md 与 DEVELOPMENT_GUIDE.md 供贡献者查阅，开始之前建议先通读。开发过程中，papermill/tests/ 下的测试套件（例如 test_engines.py 中的 TestEngineRegistration、test_iorw.py 中的 test_entrypoint_register）分别验证了 engine 与 I/O handler 的 entry point 注册逻辑，可作为自己插件行为对标的回归参考。

小结

Papermill 的扩展体系可以概括为两条清晰的插件通道：

• I/O Handler（papermill.io 组）：实现 read / write / pretty_path / listdir 四个类方法，以 URL 前缀注册，即可让 papermill 读写任意存储后端，默认回退到本地文件系统；
• Engine（papermill.engine 组）：继承 papermill.engines.Engine 并实现 execute_managed_notebook，通过 --engine <名称> 在命令行选用，可接管或增强执行逻辑（远程执行、结果后处理、重复运行模拟等）。

两者都只需在插件自己的 pyproject.toml 中声明 entry point、正常安装即可被 papermill 自动发现，完全不需要改动 papermill 本体。理解这些接口约定（以及 NotebookExecutionManager 的状态管理语义）之后，无论是接入私有存储、对接远程执行集群，还是为 notebook 注入自定义分析输出，都可以在数十分钟内落地为一个独立的可复用插件包。