python中的generator

作者: zsxwing 更新: 2012-03-17 16:42:01 发布: 2012-03-17 16:42:01

迭代器模式是一个很经典的设计模式，我们在许多语言中都可以看到它的应用，比如stl的iterator和java的Iterator接口。迭代器模式提供了一种封装遍历的功能。在python中，也有Iterator对象。但是迭代器模式需要记住当前的状态，以便返回下一数据项。这就导致在一些非线性的遍历（例如树）下，迭代器很难编写。

python提供了generator的机制来解决这个问题。generator可以帮我们把非线性的遍历封装成线性的遍历。

所谓的非线性遍历一般会涉及到递归函数，使得我们很难快速写出Iterator的实现。

例如，如果我们需要一个遍历文件夹及其所有文件并进行处理的方法。在java中，我们可能会写出这样子的方法：

1	private void walk(File parent, FileHandler fileHandler) {

2	if (parent.isDirectory()) {

3	for (File child : parent.listFiles()) {

4	walk(child, fileHandler);

5

}

6

} else {

7	fileHandler.handle(parent);

8

}

9

}

 

这是一个递归函数，需要传入一个FileHandler对象，对遍历到的文件进行处理。如果想把这个转化为迭代器模式，那么我们需要进行许多保存现场的工作。例如，下面是一个Iterator的实现。

01	public class FileIterator implements Iterator<File> {

02

03	class WalkState {

04	private WalkState(File[] currentFiles, int currentIndex) {

05	this.currentFiles = currentFiles;

06	this.currentIndex = currentIndex;

07

}

08

09	private File[] currentFiles;

10	private int currentIndex;

11

12	private boolean isEnd() {

13	return currentIndex >= currentFiles.length;

14

}

15

}

16

17	Deque<WalkState> stateStack;

18

19	public FileIterator() {

20	stateStack = new ArrayDeque<WalkState>();

21	if (file.isDirectory()) {

22	stateStack.push(new WalkState(file.listFiles(), 0));

23

} else {

24	stateStack.push(new WalkState(new File[] { file }, 0));

25

}

26

}

27

28	private File nextFile;

29

30

@Override

31	public boolean hasNext() {

32	WalkState walkState = null;

33

do {

34	if (stateStack.isEmpty()) {

35	return false;

36

}

37	walkState = stateStack.poll();

38	} while (walkState.isEnd());

39

40	nextFile = walkState.currentFiles[walkState.currentIndex++];

41	stateStack.push(walkState);

42	if (nextFile.isDirectory()) {

43	stateStack.push(new WalkState(nextFile.listFiles(), 0));

44	return hasNext();

45

} else {

46	return true;

47

}

48

}

49

50

@Override

51	public File next() {

52	return nextFile;

53

}

54

55

@Override

56	public void remove() {

57	throw new UnsupportedOperationException(

58	"Only support to traversal the dir");

59

}

60

61

}

我们需要一个栈来保存遍历的状态，这样子才能在next的时候返回正常的值。可以看到，代码很长也很容易出错。

而python的generator解决了这个问题。从而，我们可以写一个递归函数，但是提供迭代器的功能。例如，下面使用generator来遍历一个目录及其子目录的所有文件的python代码：

01

import os

02	def walk(path):

03	if os.path.isdir(path):

04	for file in os.listdir(path):

05	file_path=os.path.join(path, file);

06	for sub_file in walk(file_path):

07	yield sub_file

08

else:

09	yield path

10

11	for file in walk("."):

12	print file

这里的关键是yield，存在yield的函数称为Generator，在python中会特殊处理，。yield会把函数当前的执行状态保存，同时会将yield后面的值作为函数的返回值返回给调用者，当下次这个函数再次被调用时，会恢复函数的状态，从yield的下一个语句开始执行（不再是从函数第一行开始执行了）。这个功能帮助我们既可以写简单不易出错的递归函数，又可以提供线性的迭代器功能了（不需要手动保存状态了）。