Random Tech Thoughts

The title above is not random

Shell Like Data Processing in Python -- Using Decorators

前面的文章展示了管道的好处，以及在 Python 程序中利用管道的思想。但是前面文章里的代码还有一点缺陷，看下面的 shell 脚本和 Python 代码的比较：

find logdir -name "access-log*" | \
xargs cat | \
grep '[^-]$' | \
awk '{ total += $NF } END { print total }'

1
2
3

lines = grep('[^-]$', cat(find('logdir', 'access-log*')))
col = (line.rsplit(None, 1)[1] for line in lines)
print sum(int(c) for c in col)

Python 中 find, cat, grep 的调用用一层层的括号嵌套起来进行调用，执行顺序是从最内部的括号开始，可读性没有 shell 脚本好。可不可能在 Python 脚本中用类似 shell 脚本里的语法来提高可读性？用运算符重载就可以做到了。最初的想法来自这里，这个例子用的方法是重载或运算符 “|”。新的 Python 代码如下，我把这样的代码称为 pipe syntax (抽取最后一列并求和的代码不变。)

1

lines = find('logdir', 'access-log*') | cat | grep('[^-]$')

用运算符重载的方法把一个函数放在管道中必须自己定义一个类，在 override 的 ror() 函数中完成实际的工作。每次都要定义一个类我觉得不够方便，因此我用 decorator 来进行简化。

简单来说 decorator 通过用其他对象替换原来的函数来改变函数的行为。对 decorator 的介绍可以参考 Bruce Eckel 的两篇文章，分别介绍了无参数和有参数的 decorator 如何创建，无参数的那篇文章还介绍了 decoractor 的作用。Bruce Eckel 认为 decorator 就像是宏一样，可以改变函数的语义。他还认为 Python 的 decorator 就像 Lisp 的宏一样 powerful。对这一点我不太赞同，两者差别还是很大的，decorator 其实只是提供了在运行时动态地替换函数的功能，而 Lisp 的宏是在编译时生成代码，要说 powerful 肯定还是 Lisp 的宏更强，但 decorator 更简单而且也已经足够 powerful 了。

使用 decorator 来定义 cat 的代码如下，函数前的 “@” 是为了支持 decorator 而引入的新的语法：

1
2
3
4
5
6
7
8
9
10
11

@pipeable
def grep(iter, match):
    if callable(match):
        fun = match
    else:
        fun = re.compile(match).match
    return ifilter(fun, iter)

# Without @pipeable before the function definition of grep,
# we can use the following code to achieve the same effect.
grep = pipeable(grep)

pipeable 其实只是一个普通的 Python 对象，可以是一个函数，也可以是一个类。如果是函数，那么 grep 就是给它的参数；如果是类，grep 是给它的初始化函数的参数。pipeable(grep) 必须返回一个能够调用的对象（含有 call() 方法），除此之外没有其他要求。可见，decorator syntax 只是语法糖而已，但这个语法糖使得我们可以在函数之前加上修饰，而且从代码上一下就可以看出 grep 可以用在管道中。

下面首先描述 pipeable 修饰函数的要求和修饰后的行为，然后再看 pipeable 的实现。

可以用 pipeable 进行修饰的函数只有一个要求，第一个参数必须支持遍历操作。如果这个函数之后还有其他管道，那么函数的返回值也需要支持遍历。

被 pipeable 修饰之后，函数的行为如下：

  1. 可以像没有修饰过时一样调用，函数行为不变
  2. 调用时给定除了第一个参数以外的所有参数，当函数对象出现在 “|” 右侧时，将左侧对象作为函数的第一个参数，与之前的参数一起完成函数的调用。

用来实现 pipe syntax 的 decorator 如下（完整代码shelike.py）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

from functools import update_wrapper
class pipeable:
    def __init__(self, method):
        self.func = method
        self.args = []
        self.kwds = {}
        self.reqlen = 0
        # Since we need to allow classes to be used in pipe, there are cases that
        # method is not a function.
        if hasattr(self.func, 'func_code'):
            self.reqlen = self.func.func_code.co_argcount
        # makes the wrapper object looks like the wrapped function
        update_wrapper(self, method)

    def __call__(self, *args, **kwds):
        curlen = len(args)
        # An ugly hack to handle classes.
        if curlen == self.reqlen or 0 == self.reqlen:
            return self.func(*args, **kwds)
        elif curlen != self.reqlen - 1:
            raise TypeError('Arguments number wrong.')

        cpy = deepcopy(self)
        cpy.args = args
        cpy.kwds = kwds
        return cpy

    def __ror__(self, iter):
        return self.func(iter, *self.args, **self.kwds)

由于修改后的对象需要支持 “|” 运算符，pipeable 用类的形式实现更方便。（看了 Bruce Eckel 的文章后我也偏好用类来创建 decorator。）

  1. __init__() 接受要修饰的函数，保存起来以备后面调用，这个函数在函数被修饰时执行，实际上就是创建一个pipeable 对象，把原先函数名赋给这个对象。最后的 update_wrapper 把被修饰函数的 __name__, __doc__, __module__ 属性拷贝到创建的对象上，这样这个对象看起来就更加像原先的函数。（否则这些属性的值的都是这个对象的值。）
  2. __call__() 在被修饰函数/替换的对象被调用时执行，这里根据参数个数直接调用原函数或者把参数保存起来
  3. __ror__() 就是实现 pipe syntax 的关键，把 “|” 左边的序列和其他参数组合起来完成被修饰函数的调用。有些 ugly hack 是为了处理 class 作为初始化参数的情况

要把 shell 里的常用命令一个个用 Python 再实现一下也比较麻烦，因此我还实现了一个函数用来把 Python 里的数据转成字符串直接调用 shell 命令来处理，处理完再转成一行一行的 Python 字符串。比较下用 pipeable 和自己定义一个 class 的实现代码（现在的处理方式是把输入一次性全部转成字符串通过 OS 的管道传给外部程序，数据量大的话内存开销比较大，但我没有想到好的解决办法。）

1
2
3
4
5
6
7
8
9
10
11

@pipeable
def shell(iter, cmd=None):
    pipe = Popen(cmd, shell=True, stdin = PIPE, stdout = PIPE)
    return pipe.communicate(''.join(iter))[0].splitlines(True)

class shell:
    def __init__(self, cmd):
        self.cmd = cmd
    def __ror__(self, iter):
        pipe = Popen(self.cmd, shell=True, stdin = PIPE, stdout = PIPE)
        return pipe.communicate(''.join(iter))[0].splitlines(True)

有了这个管道，就可以把求和的任务直接用 awk 来完成了。

1
2

sumtmp = lines | shell("awk '{ total += $NF } END { print total }'") | aslist
print int(sumtmp[0])

当然，如果不想依赖外部程序，可以写一个提取字符串第 n 列/最后一列的函数，作为 tr 的参数放入管道（这个函数的名字不太好，其实就是 map，代码见shelike.py），最后转成整数以后再用 sum 求和即可。

1

print (lines | tr(last_column) | tr(int) | sum)

我很喜欢看这样的代码，有点函数式的味道，而且可读性也很好 :)

有兴趣的话 decorator module 和 PythonDecoratorLibrary 提供了更多使用 decorator 的例子，推荐。

Comments