Generatory - optymalizacja użycia pamięci w Pythonie

Załóżmy, że mamy do wykonania pozornie proste zadanie: otworzyć plik tekstowy, przetworzyć jego zawartość i zapisać wyniki do nowego pliku.

To dość powszechna operacja przy pracy z kodem, a na pierwszy rzut oka wydaje się nieskomplikowana. Przykładowa implementacja mogłaby wyglądać następująco:

source_file = "source.txt"
dest_file = "dest.txt"


with open(source_file, "r") as fsource:
    with open(dest_file, "w") as fdest:
        for source_line in fsource:
            if "@" in source_line:
                fdest.write(source_line)

Powyższy kod otwiera plik źródłowy zawierający tekst oraz plik docelowy, a następnie wyszukuje wiersze zawierające symbol @ i zapisuje je do pliku wynikowego.

Choć kod spełnia swoje zadanie, szybko staje się mało czytelny, mimo że wykonuje tylko jedną operację na danych. Gdyby liczba operacji wzrosła, jego interpretacja i modyfikacja stałyby się znacznie trudniejsze.

Dodatkowo, brak wyraźnego podziału odpowiedzialności (odczyt, przetwarzanie, zapis) utrudnia testowanie kodu oraz jego dalszą rozbudowę.

Jak to poprawić?

Lepszym podejściem jest podział kodu na funkcje, gdzie każda z nich odpowiada za jedno konkretne zadanie:

def load_file(file_name: str) -> list[str]:
    with open(file_name, "r") as fsource:
        return fsource.readlines()


def save_content(file_name: str, content: list[str]):
    with open(file_name, "w") as fdest:
        fdest.writelines(content)


def process_content(content: list[str]) -> list[str]:
    return [line for line in content if "@" in line]


if __name__ == "__main__":

    source_file = "source.txt"
    dest_file = "dest.txt"

    contents = load_file(source_file)
    processed_content = process_content(contents)
    save_content(dest_file, processed_content)
    
    

Dzięki temu kod jest bardziej przejrzysty i modułowy. Funkcja process_content może zostać łatwo rozbudowana lub podzielona na mniejsze jednostki w przypadku bardziej złożonej logiki. Ponadto, zastosowanie adnotacji typów ułatwia analizę przepływu danych (więcej na ten temat w osobnym artykule).

Nowy problem: zużycie pamięci

Choć kod został poprawiony pod względem czytelności i organizacji, pojawił się nowy problem – pamięć operacyjna.

Załóżmy, że plik źródłowy ma większy rozmiar. W obecnej implementacji cała jego zawartość jest wczytywana do pamięci.

Widzimy, że w linii 25 wzrost zużycia pamięci jest znaczący, ponad 193 MB spowodowane przez wywołanie funkcji load_file.

W przypadku jeszcze większych, plików może to prowadzić do znaczącego zużycia zasobów lub wręcz do wyczerpania dostępnej pamięci RAM.

Rozwiązanie: Generator

W naszym konkretnym przypadku nie potrzebujemy wczytywać całej zawartości pliku źródłowego do pamięci, Wystarczy odczyt jednej linii na raz:

import typing as t

def load_file(file_name: str) -> t.Generator[str, None, None]:
    with open(file_name, "r") as fsource:
        for line in fsource:
            yield line


def save_content(file_name: str, content: t.Iterable[str]):
    with open(file_name, "w") as fdest:
        for line in content:
            fdest.write(line)


def process_content(content: str) -> t.Optional[str]:
    if "@" in content:
        return content


if __name__ == "__main__":

    source_file = "source.txt"
    dest_file = "dest.txt"

    contents = load_file(source_file)
    processed_content = (line for line in contents if process_content(line) is not None)
    save_content(dest_file, processed_content)

W nowej wersji funkcja load_file nie zwraca już całej zawartości pliku, lecz generuje jego wiersze jeden po drugim, co eliminuje problem dużego zużycia pamięci. Następnie wykorzystujemy generator expression, aby filtrować linie na bieżąco i przekazać je do funkcji save_content, zapisującej przetworzone dane do pliku. 

Generator expression

Przyjrzyjmy się bliżej składni wyrażeń generatorowych (generator expression). Jest to zwięzła i czytelna forma zapisu, umożliwiająca szybkie tworzenie prostych generatorów w sposób bardziej efektywny niż tradycyjne funkcje generatorowe.

(expression for item in iterable if condition)

Co można rozpisać w następujący sposób:

def my_generator():
	for item in iterable:
		if condition:
			yield expression

Składnia wyrażeń generatorowych jest zbliżona do składni wyrażeń listowych (list comprehension). Dzięki temu osoby zaznajomione z list comprehension mogą z łatwością przyswoić zasady działania generator expressions.

Wracając do omawianego przykładu, linię:

    processed_content = (line for line in contents if process_content(line) is not None)

Moglibyśmy zapisać jako:

def filter_contents(contents: t.Iterable[str]) -> t.Generator[str, None, None]:
    for line in contents:
        if process_content(line) is not None:
            yield line

Klasyczny sposób definiowania generatora wymaga stworzenia funkcji zwracającej generator. Jednak w przypadku prostszych scenariuszy bardziej eleganckim i efektywnym rozwiązaniem jest zastosowanie skróconej składni wyrażeń generatorowych.

Podsumowanie

Oto jak wygląda zużycie pamięci w ostatecznej wersji implementacji:

Efekt jest identyczny jak w pierwszej implementacji: odczytujemy linię, przetwarzamy ją i zapisujemy do pliku wynikowego. Istnieje jednak kilka kluczowych różnic. Dzięki refaktoryzacji, udało się poprawić:

• Testowalność – poszczególne funkcje można łatwo testować niezależnie.
• Czytelność i modularność – podział na odpowiedzialne funkcje ułatwia rozwój i konserwację kodu.
• Efektywność pamięciową – umożliwia przetwarzanie nawet bardzo dużych plików, ograniczając zużycie pamięci do jednej linii naraz.

Dzięki temu kod staje się bardziej elastyczny, łatwiej skalowalny do większych zbiorów danych oraz zdolny do obsługi bardziej złożonej logiki przetwarzania bez zmniejszania czytelności, jednocześnie minimalizując ryzyko nadmiernego zużycia pamięci operacyjnej.