如何在Python中定义split_train_test？

import random

def split_train_test(data, test_ratio):
    """
    将数据集按照指定的测试集比例进行划分

    参数:
    data: 待划分的数据集，可以是列表、数组或其他可迭代对象
    test_ratio: 测试集所占的比例，取值范围为0到1之间

    返回值:
    train_set: 训练集
    test_set: 测试集
    """
    # 随机打乱数据集
    random.shuffle(data)
    # 计算测试集的大小
    test_size = int(len(data) * test_ratio)
    # 划分数据集
    test_set = data[:test_size]
    train_set = data[test_size:]
    return train_set, test_set