装饰器入门

很早之前, 前公司的Python大神就有给我们科普过装饰器的使用, 但是由于个人愚钝, 一直没有很好理解这个东西, 也没有在实际的写代码的时候用过. 因此这里把这种工作总遇到的一个例子记录一下.

首先, Python中装饰器本身的作用是在不改动原函数代码的前提下为其附加新的功能, 因为它本身其实也要求你去定义一个函数, 只不过这个函数以函数为接收对象, 也以函数为返回对象.

我本周碰到的一个问题是, 我写了一个函数去计算Bam文件中UMI序列的一些指标, 但是输入的Bam带的UMI有两种情况:

双端各N个bpUMI
单端N个bpUMI
这两种情况在从Bam文件中读取UMI的时候处理是不一样的, 单端的话要读取UMI全长, 双端的话只读UMI的半长. 原本单端的代码如下:

def single_count(file, tag, mode):
    umi_count = {
        'R1': [],
    }
    infos = load_single(file)
    for info in infos:
        if info['strand'] == 'R1':
            umi_count['R1'].append(info['umi'])
    umi_count['R1'] = Counter(umi_count['R1'])
    total_umi = 0
    read_with_umi = 0
    max_grp_size = 0
    max_umi = ''
    size_stat = Counter([1, 5, 10, 100, 1000])
    for umi in umi_count['R1']:
        if umi_count['R1'][umi] > max_grp_size:
            max_grp_size = umi_count['R1'][umi]
            max_umi = umi
        for size in size_stat: # 注意初始化的时候, 全都是1, 不是0, 所以最后要处理(-1)
            if umi_count['R1'][umi] <= size:
                size_stat[size] += 1
        total_umi += 1
        read_with_umi += umi_count['R1'][umi]
    size_stat = OrderedDict(list([("%%Group_Size_below_%s(%s)" % (size, tag), str((size_stat[size] - 1) / total_umi )) \
        for size in size_stat])) # python所以要先list化
    return total_umi, read_with_umi, max_grp_size, max_umi, size_stat

而双端读取的代码与单端的区别仅仅是infos = load_single(file)这一句需要替换成infos = load_duplex(file)而已. 为了实现这一目的其实也有很直接简单的方式, 就是更改这个函数, 加上一个参数, 然后在这一句前面加一个判断, 根据传入的参数决定调用load_duplex还是load_single. 但是这样一来会需要改动这个函数的代码, 以及调用这个函数的其他代码. 如果这样的地方一多, 其实难免会产生新的bug. 因此这次我尝试使用装饰器来解决这个问题:

def load_mod(func):
    load_single = load_duplex
    return func

@load_mod
def duplex_count(file, tag, mode):
    return single_count(file, tag, mode)

这里我新定义了两个函数, duplex_count想达到的效果是重复single_count中除了infos = load_single(file)之外的所有步骤, 然后只把调用的函数改成load_duplex. 因此可以看到这个函数的内容其实就是返回single_count的计算结果. 单这样一来并没有完成两个load函数的替换.
这里就是load_mod发挥作用的地方了: 它用来对局部命名空间内的内容进行替换, 具体是将load_single的内容以load_duplex替换. 以load_mod装饰duplex_count后, 相当于在执行函数前先把load_single的内容以load_duplex替换, 这样一来, 这个函数的执行的东西就改变了.

当然, 我这么用其实更像常见装饰器教程中的’闭包’用法, 而不是’为函数添加新功能’.

其他的应用之后再继续水.