Когда пишешь на Python, редко задумываешься, что происходит под капотом. С одной стороны, это ускоряет разработку, но, с другой, становится причиной низкой производительности и ошибок Out of memory на больших объёмах данных. Здесь мы рассмотрим несколько приёмов, как избежать подобных проблем, а в конце сравним производительность разных решений (в том числе посоревнуемся с однострочником на bash).

Начнём с типовой задачи: надо записать в файл список строк без каких-то изменений. Решение выглядит очевидным:

with open('output', 'w') as fh:
    fh.write(''.join(str_list))

Такой код первым приходит на ум, но есть проблема: перед записью мы преобразуем список в большую строку, которая займёт столько же памяти, сколько сам список. Этого легко избежать, если воспользоваться методом writelines:

with open('output', 'w') as fh:
    fh.writelines(str_list)

Теперь предположим, что после каждого элемента нужно добавить символ новой строки. Кошмарное решение выглядит так

with open('output', 'w') as fh:
    fh.write('\n'.join(str_list) + '\n')

Здесь не просто генерируется большая строка, но она ещё раз и копируется при добавлении \n в конец (строки относятся к неизменяемым типам).

Лучше взять функцию print и передать ей str_list как распакованный список аргументов:

with open('output', 'w') as fh:
    print(*str_list, sep='\n', file=fh)

Символ новой строки после последнего элемента списка будет добавлен автоматически — за это отвечает именованный аргумент end. Вообще говоря, функция позволяет добавлять любые префиксы и суффиксы:

prefix, suffix = 'foo', 'bar\n'
with open('output', 'w') as fh:
    print(prefix, end='', file=fh)
    print(*str_list, sep=suffix+prefix, end=suffix)

Ещё удобнее то, что в отличие от join, элементы списка могут быть любых типов, а не только строками:

>>> print(*range(10), sep=', ')
0, 1, 2, 3, 4, 5, 6, 7, 8, 9
>>> print('A', 10, b'B', object(), sep='-')
A-10-b'B'-<object object at 0x7f9aa43686c0>

Строка для вывода получается вызовом метода __str__ каждого элемента.

class Sep(str):
		def __init__(self, iter_):
				super().__init__()
				self.iter = iter_
        
		def __str__(self):
    		return next(self.iter)

>>> import itertools as it
>>> sep = Sep(it.cycle((', ', ', ', ',\n')))
>>> print(*range(12), sep=sep)
0, 1, 2,
3, 4, 5,
6, 7, 8,
9, 10, 11
>>> d = {'Name': 'John', 'Surname': 'Wick'}
>>> sep = Sep(it.cycle((': ', '\n')))
>>> print(*it.chain.from_iterable(d.items()), sep=sep)
Name: John
Surname: Wick

С файлами разобрались, перейдём к строкам. С ними можно делать всё то же самое с помощью объектов StringIO из стандартного модуля io. Они поддерживают те же методы, что и открытый на чтение/запись файл, и их можно передать функции print как аргумент file. Отличие только в том, что StringIO хранит данные в памяти и не имеет дескриптора файла. Данные из объекта получают одной строкой, вызывая метод getvalue(). При этом копирование данных происходит, так как строка относится к неизменяемому типу, а StringIO — к изменяемому. Также можно вернуться в начало «файла» и пробежаться по строкам:

>>> import io
>>> o = io.StringIO()
>>> o.writelines(['aina\n', 'peina\n', 'para\n'])
>>> print(*range(10), file=o)
>>> o.getvalue()
'aina\npeina\npara\n0 1 2 3 4 5 6 7 8 9\n'
>>> o.tell()
36
>>> o.seek(0)
0
>>> for line in o: print(line, end='')
... 
aina
peina
para
0 1 2 3 4 5 6 7 8 9

У StringIO есть аналог для работы с байтами BytesIO. Правда, для байтов есть, на мой взгляд, более функциональный инструмент — встроенный тип bytearray. Это изменяемая байтовая строка, а не файл. Ещё bytearray поддерживает прямую работу с памятью через memoryview. Это мы продемонстрируем в нашей финальной задаче, где сразимся со однострочником в неравной битве. Итак, нужно прочитать несколько файлов, заменить одну последовательность байт на другую и записать результат в сокет. Решение на bash'e:

$ cat file.1 file.2 | sed 's/something/something_new/g' | nc -q0 ::1 8000

На 8000 порту в это время запущен netcat, отправляющий всё в /dev/null:

$ nc -lp 8000 -6 -k > /dev/null

Сначала напишем прямое решение без заморочек, не думая, что и зачем там копируется.

import sys
import socket


OLD = b'alert'
NEW = b'ok'


def main():
    content_list = []
    for filename in sys.argv[1:]:
        with open(filename, 'rb') as fhandler:
            content_list.append(fhandler.read())
    content = b''.join(content_list)
    content = content.replace(OLD, NEW)
    sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
    sock.connect(('::1', 8000))
    sock.sendall(content)
    sock.close()


if __name__ == '__main__':
    main()

import os
import io
import sys
import socket
import itertools as it


OLD = b'alert'
NEW = b'ok'


def read_data(filenames):
    sizes = [os.path.getsize(name) for name in filenames]
    # выделяем блок памяти, куда прочитаем содержимое файлов
    data = bytearray(sum(sizes))
    # memoryview даёт доступ непосредственно к выделенной памяти
    # здесь нужен для записи данных по смещению
    mview = memoryview(data)
    # список со смещениями файлов внутри нашего буффера
    offsets = [0]
    offsets.extend(it.accumulate(sizes))
    for name, offset in zip(filenames, offsets):
        with open(name, 'rb') as fhandler:
            # readinto читает данные непосредственно
            # в выделенный буффер без создания нового объекта
            fhandler.readinto(mview[offset:])
    return data


def sub(data, output, old, new):
    prev_offset = 0
    old_len = len(old)
    # здесь memorview нужен для получения срезов данных
    # без копирования
    mview = memoryview(data)
    while True:
        try:
            # ищем смещение заменяемой строки, начиная 
            # с предыдущего вхождения
            offset = data.index(old, prev_offset)
        except ValueError:
            # дописываем хвост с данными без old
            output.write(mview[prev_offset:])
            break
        # записываем блок между соседними позициями old
        output.write(mview[prev_offset:offset])
        # вместо самой строки old пишем new
        output.write(new)
        prev_offset = offset + old_len


def main():
    filenames = sys.argv[1:]
    data = read_data(filenames)
    sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
    sock.connect(('::1', 8000))
    # будем писать в сокет, как в файл
    output = sock.makefile(mode='wb')
    sub(data, output, OLD, NEW)


if __name__ == '__main__':
    main()

В первый решении копирование данных происходит дважды: при выполнении join и replace. Во втором — ни разу. Проверим, как это отразится на времени работы, запустим каждый вариант на двух файлах по полтора гигабайта:

$ ls -l file*
-rw-rw-r-- 1 znahar znahar 1521240266 фев  5 16:04 file1
-rw-rw-r-- 1 znahar znahar 1521238713 фев  5 16:53 file2
$ /usr/bin/time -f "\t%Es time,\t%MK memory,\t%P CPU" ./catsednc.py file1 file2
	0:02.89s time,	2980796K memory,	94% CPU
$ /usr/bin/time -f "\t%Es time,\t%MK memory,\t%P CPU" ./catsednc2.py file1 file2
	0:06.85s time,	8923064K memory,	92% CPU
$ /usr/bin/time -f "\t%Es time,\t%MK memory,\t%P CPU" cat file1 file2 | sed 's/alert/ok/g' | nc -q0 ::1 8000
	0:07.34s time,	2256K memory,	25% CPU

Какие отсюда следуют выводы:

• оптимизированный скрипт более чем в два раза быстрее других вариантов, его потребление памяти ожидаемо равно суммарному размеру файлов;

• простой скрипт, дважды копирующий данные, потребляет в три раза больше памяти;

• однострочник читает данные построчно и использует конвейеры, через которые передаются блоки данных фиксированного размера. Это радикально сокращает потребление памяти, которое больше не зависит от размера файлов, но негативно сказывается на времени выполнения. Кроме того, нагрузка на процессор гораздо меньше, так как утилиты написаны на C.

На этом всё. Надеюсь, информация будет полезной.