没想到有一天写 python 的时候也会想着如何去节省内存。平时写 python 的时候根本不会关注这些，变量什么的直接创建和使用就完了，也不用考虑内存的释放，反正有垃圾回收机制。只不过这次数据量过大，debug 的时候发现内存一直在申请，导致系统彻底的卡死。

可能也是从事算法的优化工作养成了职业病，每次写代码的时候都会想，这些代码消耗的时间怎么样，占用的空间怎么样，数据结构是否可以继续优化，这些逻辑有没有更优雅的写法。

注：本文程序中使用 psutil 库来监测进程使用的内存大小，需要 pip install psutil一下。

背景

需要解析一个很大的日志文件，日志文件中含有一些无用的信息，像下面这样：

有用信息1
无用信息1
有用信息2
有用信息3
无用信息2
...
有用信息N

解析文件的时候，需要从文件中解析并提取出有用的信息，存入一个对象中，完成后续的处理。
但是呢，对于某些特殊的任务和需求，发现文件只解析一次是不行的，也就是需要对文件进行二次解析。

所以为了避免重复的解析文件，在第一次文件解析完毕后，直接把有用的核心信息序列化出去，这样二次解析的话就不用重新读取源文件在解析，直接读取序列化后的核心数据就好了。

序列化导出

最开始的方案是使用一个 list 持续追加解析得到的核心数据，文件解析完毕后把这个很大的 list 序列化出去。监测到进程占用的内存大小为：700MB。

import random
import pickle
import time
import psutil
import os

data = []

for i in range(10000000):
    data.append(str(random.randint(10000, 109070987)))

with open("data.pkl", "wb") as f:
    pickle.dump(data, f)

# 获取当前 Python 进程占用的内存
memory_info = process.memory_info()

# 打印占用的内存大小，rss 单位为字节
print(memory_info.rss / 1024 / 1024, "MB")

而如果使用序列化追加的方式，仅用 15MB，耗时增加 2s，毕竟每次序列化的时候都需要打开文件并在末尾追加内容：

1
2
3

with open("data.pkl", "ab") as f:
    for i in range(10000000):
        pickle.dump(str(random.randint(10000, 109070987)), f)

这里可以设置一个 buffer 进行优化，buffer 达到一定大小后在统一序列化出去。

class SeriesModel:
    def __init__(self) -> None:
        self._buf = []

    def series(self, stack, finish=False):
        self._buf.append(stack)
        if 100 < len(self._buf) or finish is True:
            with open(config.SERIES_PATH, "ab") as f:
                for item in self._buf:
                    pickle.dump(item, f)
                self._buf = []

序列化读入

在二次解析的时候，需要把序列化的数据 load 进来。如果加载序列化的文件并且直接处理数据，同样需要使用 700MB 的内存。这种一次性创建所有元素的行为是没有必要的。

with open("data.pkl", "rb") as f:
    data = pickle.load(f)

    for i in data:
        i += " "

可以使用惰性计算来解决这一问题，只有在真正需要这个变量的时候才去创建，而不是一开始就创建所有的变量。考虑到生成器表达式的局限性，我们直接使用 yield 关键字创建一个生成器函数。

yield 语句类似 return 会返回一个值，但它会记住这个返回的位置，下次迭代的时候就从这个位置继续执行，返回下一个元素。这样就消耗内存 15MB。

def read(file):
    with open(file, "rb") as f:
        data = pickle.load(f)
        for i in data:
            yield i

# data 是生成器
data = read("data.pkl")
for i in data:
    i += " "

引申

任何一个生成器都会定义一个名为 __next__ 的方法，这个方法要在最后一个元素之后需抛出 StopIteration 异常。next() 函数的本质就是调用对象的 __next__()。这个方法要么返回迭代的下一项，要么引起结束迭代的异常 StopIteration，下面的示例揭示了生成器的本质。

class FibGenerator():
    def __init__(self, n):
        self.__n = n

        self.__s0 = 0
        self.__s1 = 1
        self.__count = 0

    def __next__(self):  # 用于内建函数 next()
        if self.__count < self.__n:
            ret = self.__s0
            self.__s0, self.__s1 = self.__s1, (self.__s0 + self.__s1)
            self.__count += 1
            return ret
        else:
            raise StopIteration

    def __iter__(self):  # 用于 for 循环语句
        return self

fg = FibGenerator(5)
print(type(fg))
print(isinstance(fg, Iterable))
for i in fg:
    print(i, end=' ')

>>>
<class '__main__.FibGenerator'>
True
0 1 1 2 3

示例中如果没有定义 __iter__() 方法则只能使用 next() 函数进行迭代，当它定义后，就可以使用 for 和 in 语句访问了，同时定义了这两种方法的对象称为迭代器。生成器表达式和生成器函数产生生成器时，会自动生成名为 __iter__ 和 __next__ 的方法，所以生成器也是一种迭代器。

参考链接

https://pythonhowto.readthedocs.io/zh-cn/latest/iterator.html

Just for Life.

在 python 操作大文件时节省内存

背景

序列化导出

序列化读入

引申

参考链接