Сообщить об ошибке.

Паттерны использования io.StringIO() в Python

Содержание:

io.StringIO - это "файл в памяти" для строк (str). Почти как BytesIO, но уже текстовый, с теми же методами, что у обычного текстового файла (read, write, readline, seek, итерация по строкам).

Разберёмся по паттернам io.StringIO: от самых частых до более хитрых.

Временный текстовый файл "внутри памяти"

Самый базовый паттерн: нужно что-то, что умеет писать/читать строки, как файл - используем StringIO.

from io import StringIO

f = StringIO()
f.write("hello ")
f.write("world\n")
f.seek(0)

print(f.read())      # "hello world\n"

Когда использовать:

  • вместо временного файла на диске;
  • когда данные маленькие/средние по объёму;
  • когда API ожидает file-like объект с текстом.

Подмена файла в тестах

Классика: функция читает/пишет в файл - в тестах подсовываем StringIO.

from io import StringIO

def process_file(f):
    # f - любой текстовый поток: open(...), StringIO и т.п.
    lines = [line.strip() for line in f]
    return ",".join(lines)

def test_process_file():
    fake = StringIO("a\nb\nc\n")
    assert process_file(fake) == "a,b,c"

Плюсы:

  • не трогаем файловую систему;
  • проще проверять содержимое (строки, не байты);
  • легко симулировать ошибки (например, сделать кастомный StringIO, который кидает исключение после N read).

Захват вывода (stdout/stderr)

Паттерн: временно заменить sys.stdout/sys.stderr на StringIO, чтобы поймать вывод.

import sys
from io import StringIO

buf = StringIO()
old_stdout = sys.stdout
sys.stdout = buf

try:
    print("hello")
    print("world")
finally:
    sys.stdout = old_stdout

captured = buf.getvalue()
print("Captured:", repr(captured))  # "hello\nworld\n"

Частый сценарий:

  • тесты для CLI;
  • логирование "что бы напечаталось";
  • прогон чужого кода, чтобы не шумел в консоль.

Генерация больших строк без кучи конкатенаций

Паттерн: накапливать куски текста через write, а в конце делать getvalue().

from io import StringIO

buf = StringIO()

for i in range(5):
    buf.write(f"line {i}\n")

result = buf.getvalue()
print(result)

Почему так: конкатенация "строка" + "строка" в цикле создаёт много временных объектов; StringIO под капотом делает это более оптимально.

Адаптер для API, которое ожидает текстовый файл

Множество функций/библиотек принимают параметр типа file-like object, например:

  • CSV (csv.reader / csv.writer),
  • configparser,
  • некоторые XML/JSON-парсеры (или наоборот, генераторы).

Пример: CSV из строки

import csv
from io import StringIO

csv_text = "name,age\nAlice,30\nBob,25\n"

f = StringIO(csv_text)
reader = csv.DictReader(f)

for row in reader:
    print(row)
    # {'name': 'Alice', 'age': '30'}, ...

Пример: сгенерировать CSV в строку

f = StringIO()
writer = csv.writer(f)
writer.writerow(["name", "age"])
writer.writerow(["Alice", 30])
writer.writerow(["Bob", 25])

csv_text = f.getvalue()
print(csv_text)

Временный лог / "текстовый буфер" для последующей отправки

Хочется "писать лог" в стильном API (write, print(..., file=...)), а потом всё это отправить письмом/по сети/записать в одно место.

from io import StringIO
from datetime import datetime

log = StringIO()

def log_msg(msg):
    ts = datetime.now().isoformat(timespec="seconds")
    print(f"[{ts}] {msg}", file=log)

log_msg("start")
log_msg("doing stuff")
log_msg("finish")

## потом, например, отправляем содержимое
body = log.getvalue()
print("EMAIL BODY:\n", body)

"Псевдо-файл" для функций, которые печатают в файл

Любой print умеет выводить в произвольный поток через file=.

from io import StringIO

buf = StringIO()
print("hello", "world", file=buf)
print("another line", file=buf)

text = buf.getvalue()
print(repr(text))  # 'hello world\nanother line\n'

Паттерн:

  • есть функция, которая "жёстко" использует print(..., file=...);
  • вы передаёте туда StringIO и получаете результат как строку.

Сброс/повторное использование StringIO

Часто удобно не создавать каждый раз новый объект, а "очищать" и использовать снова:

from io import StringIO

buf = StringIO()
buf.write("old data")
buf.seek(0)
buf.truncate(0)   # стереть содержимое
buf.write("new data")

buf.seek(0)
print(buf.read())  # "new data"

Паттерн:

  • seek(0) + truncate(0) - очистить;
  • seek(0) - перемотать для чтения.

Построчный разбор текста

Если есть большая строка, но её удобно парсить как поток:

from io import StringIO

data = "line1\nline2\nline3\n"
f = StringIO(data)

for line in f:
    print(">", line.strip())

Паттерн: StringIO как удобный итератор по строкам вместо splitlines() - особенно, если нужно сложное чтение (read, readline, смешанное с seek).

Отличия от BytesIO и когда какой использовать

Кратко:

  • StringIO:

    • работает только со строками (str);
    • сам не знает про кодировку - это просто абстракция уровня текста;
    • идеален для тестов, логов, CSV/JSON/XML и всего, что уже в str.

  • BytesIO:

    • работает с bytes;

    • нужен, когда:

      • ты общаешься с байтовыми API (архивы, картинки, сети),
      • хочешь потом обернуть в TextIOWrapper с явной кодировкой.

Частый паттерн:

  • на нижнем уровне - BytesIO (байты),
  • поверх - TextIOWrapper => результат читаешь/пишешь как str.

А StringIO - когда байты вообще не нужны, всё уже текст.

Примеры с StringIO

Кастомный StringIO-подобный для особых случаев

Иногда нужно делать что-то ещё при записи/чтении:

  • считать символы;
  • автоматически добавлять префиксы;
  • фильтровать строки.

Часто удобно наследоваться от StringIO:

from io import StringIO

class CountingStringIO(StringIO):
    @property
    def length(self):
        pos = self.tell()
        self.seek(0, 2)  # конец
        end = self.tell()
        self.seek(pos)
        return end

buf = CountingStringIO()
buf.write("hello")
buf.write(" мир")
print(buf.length)   # 11

Или сделать декоратор, который оборачивает обычный StringIO/TextIOBase, но это уже другая история.

Пример подмены sys.stdin вручную

Допустим, есть функция, которая читает из input():

def ask_name():
    name = input("Your name: ")
    print(f"Hello, {name}!")
    return name

В тесте можно подменить sys.stdin на StringIO:

import sys
from io import StringIO

def test_ask_name():
    fake_input = StringIO("Alice\n")  # то, что 'напечатает' пользователь

    old_stdin = sys.stdin
    sys.stdin = fake_input
    try:
        name = ask_name()
    finally:
        sys.stdin = old_stdin  # обязательно восстановить!

    assert name == "Alice"

Что тут происходит:

  • input() читает из sys.stdin;
  • мы подсовываем туда StringIO, в котором лежит строка "Alice\n";
  • функция думает, что это настоящий stdin, и работает как обычно.

Пример контекстного менеджера для аккуратной подмены

Чтобы не забывать восстанавливать sys.stdin, удобно сделать маленький контекстный менеджер:

import sys
from io import StringIO
from contextlib import contextmanager

@contextmanager
def fake_stdin(text: str):
    """
    Временно подменяет sys.stdin на StringIO(text).
    Возвращает объект StringIO (чтобы можно было ещё и читать, если нужно).
    """
    old_stdin = sys.stdin
    buf = StringIO(text)
    sys.stdin = buf
    try:
        yield buf
    finally:
        sys.stdin = old_stdin

Использование:

def ask_two_numbers():
    a = int(input("A: "))
    b = int(input("B: "))
    return a + b

def test_ask_two_numbers():
    # имитируем ввод: "2\n3\n"
    with fake_stdin("2\n3\n"):
        s = ask_two_numbers()
    assert s == 5

Пример нескольких вызовов input() подряд

Важно: всё, что input() читает, идёт по строкам - до \n. Поэтому для нескольких вызовов input() просто кладём несколько строк в StringIO:

from io import StringIO
import sys

def dialog():
    name = input("Name: ")
    city = input("City: ")
    return f"{name} from {city}"

def test_dialog():
    fake_input = StringIO("Alice\nRiga\n")

    old_stdin = sys.stdin
    sys.stdin = fake_input
    try:
        res = dialog()
    finally:
        sys.stdin = old_stdin

    assert res == "Alice from Riga"

Пример подмены stdin для функции, которая сама читает sys.stdin напрямую

Если функция не использует input(), а сама делает sys.stdin.read() / for line in sys.stdin, паттерн тот же:

import sys
from io import StringIO

def read_all_stdin_upper():
    data = sys.stdin.read()
    return data.upper()

def test_read_all_stdin_upper():
    fake_stdin = StringIO("hello\nworld")
    old_stdin = sys.stdin
    sys.stdin = fake_stdin
    try:
        res = read_all_stdin_upper()
    finally:
        sys.stdin = old_stdin

    assert res == "HELLO\nWORLD"

Заметка для pytest

С pytest обычно делают так:

def test_ask_name(monkeypatch):
    from io import StringIO

    fake_input = StringIO("Alice\n")
    monkeypatch.setattr("sys.stdin", fake_input)

    name = ask_name()
    assert name == "Alice"

Но идея всё та же: StringIO => в sys.stdin => функция думает, что общается с пользователем.