基于 entry points 扩展 papermill：自定义 I/O 处理器与执行引擎完整指南

本篇指南以 papermill 的扩展机制为专题，讲解如何利用 Python 打包生态中的 entry points 机制，为 papermill 开发全新的输入/输出（I/O）处理器与自定义执行引擎。读完本文，你将掌握 I/O 处理器四个接口方法的实现约定、在 pyproject.toml 中注册 papermill.io 与 papermill.engine 两组 entry points 的标准写法，并能够从零实现一个支持 sftp:// 路径的读写插件和一个为每个代码单元注入执行耗时输出的自定义引擎（通过 --engine 参数一键切换）。

什么是 entry points：papermill 插件化的基石

Python 官方打包规范对 entry points 的定义是：一种让已安装的发行包（distribution）对外「广告」自身组件、供其他代码发现和使用的机制。典型场景包括：

• 发行包可以声明 console_scripts entry points，每个都指向一个函数；当 pip 等安装器安装该发行包时，会为每个 entry point 生成对应的命令行包装器（例如 papermill 自身的命令行入口 scripts.papermill = "papermill.__main__:papermill"，见仓库根目录 pyproject.toml）。
• 应用程序可以利用 entry points 加载插件，例如语法高亮工具 Pygments 可以从独立安装的第三方包中获取额外的 lexer 与 style。

papermill 正是基于这一机制构建插件体系的：运行时它会扫描两类 entry points——实现输入/输出（I/O）处理的 papermill.io 组，以及实现 notebook 执行逻辑的 papermill.engine 组。这意味着你不需要修改 papermill 源码，只需编写一个独立的小包并声明对应的 entry points，安装后即可被 papermill 自动发现与调用。

从仓库源码可以印证这一点。papermill 在模块加载时即完成插件扫描：

• 在 papermill/iorw.py 中，模块底部实例化全局单例 papermill_io = PapermillIO()，注册内置的 local、s3://、adl://、abs://、http://、https://、gs://、hdfs://、GitHub、-（标准流）等处理器后，立即调用 papermill_io.register_entry_points()；
• 在 papermill/engines.py 中，papermill_engines 单例先注册默认引擎 NBClientEngine（同时以 None 和 'nbclient' 两个名字注册），随后同样调用 register_entry_points()。

两个 register_entry_points 的实现完全对称，都是遍历 entrypoints.get_group_all(...) 返回的 entry point 对象，将其 name 与 load() 出的类注册进内部注册表，分别见 papermill/iorw.py 与 papermill/engines.py。因此，扩展 papermill 的完整链路是：编写插件类 → 在插件包的 pyproject.toml 中声明 entry points → 安装插件包 → papermill 启动时自动注册。

papermill 的可扩展点：I/O 处理器与执行引擎

当 papermill 执行一个 notebook 时，整体流程可概括为四步（参见 docs/extending-overview.rst）：

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

通过 entry points，你可以在第 1、3、4 步注入自定义实现：I/O 处理器接管第 1 步（读）与第 4 步（写），执行引擎接管第 3 步（执行）。也就是说，papermill 的扩展面就这两类：「数据从哪里来、到哪里去」由 I/O 处理器决定，「notebook 怎么跑」由执行引擎决定。

其中 I/O 处理器让 papermill 不局限于本地文件系统，也能访问 Amazon S3、Azure Data Lake、Azure Blob Storage、Google Cloud Storage、HDFS、GitHub 乃至任意远程服务；执行引擎则决定了 notebook 是在本机执行、远程服务器执行，还是被多次重复执行以模拟不同条件，抑或在执行后对结果做额外的后处理。

开发自定义 I/O 处理器

接口约定：四个必需方法

编写一个自定义 I/O 处理器，本质上就是编写一个包含以下四个类方法的类：

方法	签名	职责
read	CustomIO.read(file_path)	返回文件内容
write	CustomIO.write(file_content, file_path)	将内容写入目标路径，无返回值
pretty_path	CustomIO.pretty_path(path)	返回美化后的路径（用于日志展示）
listdir	CustomIO.listdir(path)	返回该路径下的路径列表

注意：如果你的处理器只面向单一用途（例如仅用于发布的平台只写不读），那么不需要支持的方法也必须实现，并在被调用时抛出异常（如 raise NotImplementedError），以符合接口契约。

这一约定与 papermill 内置处理器的实现完全一致。例如内置的 S3Handler（papermill/iorw.py）四个方法俱全：read 拼接 S3().read(path) 的行、write 调用 S3().cp_string(buf, path)、listdir 委托给 S3 客户端、pretty_path 原样返回路径。而只读性质的处理器则会在不支持的方法上抛异常，例如 GithubHandler 的 write 与 listdir 直接抛出 PapermillException（papermill/iorw.py），NotebookNodeHandler 的 write 同样如此（papermill/iorw.py）——这正印证了文档中「不打算支持的方法也要实现并抛异常」的建议。

让 papermill 发现你的处理器：注册 entry points

开发完成后，需要在插件包的 pyproject.toml 中声明 [project.entry-points."papermill.io"] 段，例如：

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

这段配置的含义是：当某个文件路径以 sftp:// 开头时，papermill 应当使用 papermill_sftp.SFTPHandler 类来处理该路径的读写。等号左边是路径前缀，等号右边是类的导入路径（模块:类名）。

传统上，papermill I/O 处理器的 entry point 名称看起来都像 URL 前缀。例如 Amazon Web Services 的 S3 处理器注册在 s3:// 下，凡是路径以 s3:// 开头的读写都会交给它处理。papermill 内置处理器在 papermill/iorw.py 中的注册方式与此一一对应："s3://"、"adl://"、"abs://"、"gs://"、"hdfs://"、"http://"、"https://"、"http://github.com/" 等。

底层匹配原理（来自源码）：PapermillIO.get_handler(path)（papermill/iorw.py）按注册顺序遍历内部 self._handlers 列表，逐个判断 path.startswith(scheme)，命中即返回对应处理器；所有带前缀的处理器都不匹配时，回退到 local 处理器；如果连 local 处理器都不存在，则抛出 PapermillException("Could not find a registered schema handler for: ...")。此外 register() 使用 insert(0, ...) 维护 LIFO 顺序（papermill/iorw.py），即后注册的优先级更高，这与测试 test_register_ordering 验证的行为一致（papermill/tests/test_iorw.py）。另外，若路径不以 .ipynb/.json 结尾，get_handler 还会发出扩展名相关的 UserWarning（papermill/iorw.py）。

实战示例：完整的 sftp I/O 处理器

下面通过一个完整示例演示如何编写一个「从 SFTP 服务器读取 notebook、并把执行结果写回 SFTP」的处理器，使我们可以直接执行：

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

I/O 处理器定义在 src/papermill_sftp/__init__.py 中。设计思路：读取时先把远程文件下载到本地临时目录再读入；写入时先写本地临时文件再上传；pretty_path 不做路径改写直接返回；listdir 暂不实现：

import os
import pysftp

sftp_username = os.getenv('SFTP_USERNAME')
sftp_password = os.getenv('SFTP_PASSWORD')

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 完整内容如下（需注意 cnopts 需按你的 pysftp 使用场景自行构造，例如关闭主机密钥校验时使用 pysftp.CnOpts() 并设置 cnopts.hostkeys = None；示例代码仅演示结构）：

[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 完成读写。读者可以对照 papermill/iorw.py 中内置处理器（如 S3Handler、GCSHandler）的写法，体会「薄封装、四方法、委托给底层客户端」的通用模式。

开发自定义执行引擎

引擎的职责与接口

papermill 引擎是一个能够运行（执行）notebook 的 Python 对象。papermill 的默认实现接收一个 notebook 对象，并在本机执行它。通过编写自定义引擎，你可以让执行发生在远程服务器上，或者对执行后的 notebook 施加额外的后处理。

自定义引擎必须继承自 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 对象，kernel_name 是字符串。在你的自定义类中，需要真正实现 notebook 的执行逻辑——例如插入在远程服务器上执行 notebook 的代码，或多次执行 notebook 以模拟不同条件。

源码级原理：Engine.execute_notebook（papermill/engines.py）是所有引擎的统一入口，它先把 notebook 包装进 NotebookExecutionManager（nb_man），依次调用 nb_man.notebook_start()、cls.execute_managed_notebook(nb_man, kernel_name, log_output=log_output, **kwargs)，最后在 finally 中确保 nb_man.notebook_complete() 被调用。NotebookExecutionManager（papermill/engines.py）封装了执行状态跟踪：它为每个单元维护 pending / running / completed / failed 四种状态，并负责将 start_time、end_time、duration、exception 等执行计时与异常信息写入 notebook 元数据（见 notebook_start、cell_start、cell_complete 等回调）。因此你的自定义引擎可以完全复用这套状态机——这正是下文 timing 引擎能够读取 cell.metadata.papermill.start_time / end_time 的原因。基类的 execute_managed_notebook 默认抛出 NotImplementedError（papermill/engines.py），强制子类给出实现。papermill 自带的默认引擎是 NBClientEngine（papermill/engines.py），它基于 nbclient 在本地启动内核并逐单元执行。

实战示例：注入单元执行耗时的 timing 引擎

下面的示例项目实现了一个自定义引擎：在每个代码单元的输出之前，额外注入一条「该单元执行耗时」的输出信息。

项目结构：

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

src/papermill_timing/__init__.py 实现引擎逻辑。由于 papermill 已经将执行计时信息存放在 metadata 中，我们可以直接复用默认引擎完成执行，再用 nbformat 库构造 notebook 节点对象来插入新输出：

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(...) 先复用默认的 nbclient 引擎完成真正的执行（此时每个单元的 cell.metadata.papermill 已由 NotebookExecutionManager 写入 start_time 与 end_time），随后遍历 nb_man.nb.cells，对每个已执行（execution_count is not None）的代码单元计算耗时，用 nbformat.v4.new_output("display_data", ...) 构造一个文本展示单元，并把它插入到该单元原有 outputs 的最前面。

让 papermill 发现你的引擎：注册与命令行使用

自定义引擎需要在 pyproject.toml 中以 papermill.engine 前缀声明为 entry points，指向刚刚实现的类。例如，把名为 TimingEngine 的引擎放在名为 papermill_timing 的包中：

[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_timing 包的引擎。这一机制与 CLI 的实现直接对应：papermill 命令的 --engine 选项（papermill/cli.py）会把引擎名传入 execute_notebook(..., engine_name=engine, ...)（papermill/cli.py），最终由 PapermillEngines.execute_notebook_with_engine（papermill/engines.py）从注册表按名字取出引擎并执行；若传入的引擎名未注册，get_engine 会抛出 PapermillException("No engine named '...' found")（papermill/engines.py）。Python API 使用者则可以在调用 papermill.execute_notebook(...) 时传入 engine_name="timer_engine"（见 papermill/execute.py 中通过 papermill_engines.execute_notebook_with_engine(engine_name, ...) 分发）。

下图对比了同一份 notebook 的两种执行结果：左侧是用新自定义引擎执行的 notebook，右侧是用标准 papermill 引擎执行的 notebook。可以看到，自定义引擎在每个代码单元下都注入了「Execution took ... seconds」的额外输出。

扩展机制的测试验证与调试

papermill 仓库自带针对插件机制的单元测试，可作为你开发插件时的行为规范参考：

• papermill/tests/test_iorw.py 的 test_entrypoint_register 通过 mock entrypoints.get_group_all 验证：register_entry_points() 必须以 "papermill.io" 为分组名调用，且 entry point 的 name 被注册为路径前缀、load() 的返回值被注册为处理器——随后 get_handler("fake-from-entry-point://") 能取回该处理器。test_register_ordering（papermill/tests/test_iorw.py）则验证了前缀匹配的优先级规则。
• papermill/tests/test_engines.py 的 test_registering_entry_points 以相同方式验证引擎 entry points 以 "papermill.engine" 分组注册，且注册后可通过 get_engine("fake-engine") 取回；test_getting（papermill/tests/test_engines.py）则验证了未注册引擎名会抛出 PapermillException。

调试插件时，建议按以下顺序排查：

1. 确认插件包已安装：运行 pip show <你的包名>，或检查 pip list 中是否包含该包；
2. 确认 entry points 已注册：Python 中执行 import entrypoints; entrypoints.get_group_all("papermill.io")（或 "papermill.engine"）查看返回结果，对比 papermill/iorw.py 与 papermill/engines.py 中 papermill 的扫描逻辑；
3. 核对路径前缀：I/O 处理器的前缀匹配是 path.startswith(scheme)（papermill/iorw.py），注意大小写、协议名与 :// 后缀是否与你的路径完全一致；
4. 检查依赖声明：若插件依赖 pysftp、nbformat 等库，需在 dependencies 中声明（参考上文两个 pyproject.toml 示例，这也与 papermill 本身声明 entrypoints 依赖的方式一致，见根目录 pyproject.toml）。

最佳实践与边界

• 只读/只写平台的接口完整性：I/O 处理器必须完整实现四个方法；不支持的读写操作应抛出明确的异常（NotImplementedError 或 PapermillException），参考 papermill/iorw.py 中 GithubHandler 的做法。
• 前缀命名的规范性：遵循「URL 前缀」惯例命名 I/O 处理器 entry points（如 s3://、sftp://），避免与内置前缀（local、s3://、adl://、abs://、http://、https://、gs://、hdfs://、GitHub、-，见 papermill/iorw.py）冲突；引擎 entry point 名使用语义化的小写标识（如 timer_engine）。
• 复用默认引擎而非从零实现：绝大多数后处理型引擎都应继承 NBClientEngine 并调用 super().execute_managed_notebook(...) 完成真实执行（papermill/engines.py），只在执行前后追加自定义逻辑；执行状态与计时元数据由 NotebookExecutionManager 统一维护（papermill/engines.py），无需重复造轮子。
• 执行环境的依赖前置：sftp 示例依赖环境变量 SFTP_USERNAME / SFTP_PASSWORD 提供凭据，且代码中使用的 cnopts、urllib、tempfile、pathlib 等需要你在真实插件中补齐 import 与配置，示例旨在演示结构与数据流，而非开箱即用的生产代码。

总结

通过 entry points，papermill 的 I/O 层与执行层都是完全可插拔的：papermill.io 分组允许你为任意路径前缀（sftp://、自定义协议等）注册读写处理器，papermill.engine 分组允许你注册可在命令行用 --engine <名字> 切换的执行引擎。本文的两个完整示例——SFTP I/O 处理器与执行耗时注入引擎——覆盖了插件从「实现类」到「pyproject.toml 注册」再到「被 papermill 自动发现与调用」的全过程。结合 papermill/iorw.py、papermill/engines.py 与仓库测试 papermill/tests/test_iorw.py、papermill/tests/test_engines.py，你可以进一步深入理解 papermill 的插件加载细节，并以此为模板开发自己的远程执行、多云存储或结果增强型扩展。