如何对从API取数的dlt source进行采样，无需修改源码？

不用每次修改source代码，你可以试试这几种灵活的方法：

1. 调用source时用生成器切片限制输出 #

直接在调用source的地方用itertools.islice截断生成器，完全不用改动source内部代码：

import itertools
from your_module import get_page  # 导入你的source生成器

# 只取前10页结果
sample_data = itertools.islice(get_page(endpoint, headers, params), 10)
for page in sample_data:
    # 处理样本数据
    print(page)

如果你的get_page每次yield的是单条记录，这里的10就对应限制记录数；如果是整页数据，就对应限制页数，刚好匹配你原来的计数器逻辑。

2. 给source添加可配置参数 #

给get_page加个可选的max_pages参数，默认设为None（生产模式不限制），测试时传具体数值即可：

def get_page(endpoint, headers, params, max_pages=None):
    res = requests.get(endpoint, headers=headers, params=params).json()
    count = 0
    while res is not None:
        yield res["result"]
        count += 1
        # 仅当测试传值时触发限制
        if max_pages is not None and count >= max_pages:
            return
        has_more = res.get("paging", {}).get("next", None)
        if has_more:
            next_url = has_more["link"]
            res = requests.get(next_url, headers=headers).json()
        else:
            res = None

• 生产环境调用：get_page(endpoint, headers, params)（无限制拉取）
• 测试时调用：get_page(endpoint, headers, params, max_pages=10)（仅取10页）

3. 用装饰器包装生成器限制调用次数 #

写一个通用的装饰器，用来限制生成器的输出次数，完全不用修改原source代码：

def limit_generator(max_count):
    def decorator(gen_func):
        def wrapper(*args, **kwargs):
            gen = gen_func(*args, **kwargs)
            count = 0
            for item in gen:
                if count >= max_count:
                    return
                yield item
                count += 1
        return wrapper
    return decorator

# 测试时临时包装你的source
from your_module import get_page
limited_get_page = limit_generator(10)(get_page)

# 使用包装后的生成器获取样本
for page in limited_get_page(endpoint, headers, params):
    print(page)

这种方式可以快速给任意生成器添加限制，测试完移除包装即可，对原代码零侵入。

4. 通过环境变量控制 #

在source里读取环境变量，测试时设置环境变量启用限制，生产环境不设置则正常拉取：

import os

def get_page(endpoint, headers, params):
    res = requests.get(endpoint, headers=headers, params=params).json()
    # 从环境变量获取最大页数，默认None（不限制）
    max_pages = os.getenv("DLT_TEST_MAX_PAGES")
    max_pages = int(max_pages) if max_pages is not None else None
    count = 0
    while res is not None:
        yield res["result"]
        count += 1
        if max_pages is not None and count >= max_pages:
            return
        has_more = res.get("paging", {}).get("next", None)
        if has_more:
            next_url = has_more["link"]
            res = requests.get(next_url, headers=headers).json()
        else:
            res = None

测试前在终端设置环境变量：

# Linux/macOS
export DLT_TEST_MAX_PAGES=10

# Windows cmd
set DLT_TEST_MAX_PAGES=10

# Windows PowerShell
$env:DLT_TEST_MAX_PAGES=10

内容的提问来源于stack exchange，提问作者Violetta Mishechkina