for i in zipfile:

... Petr Blahoš, 12. 10. 2018 Python

Nedávno jsem řešil problém procházení adresářové struktury a nalezení a zpracování určitého typu souborů. S drobnou komplikací: Některé soubory mohou být zip archívy, jiné soubory mohou být komprimované jako gzip. Nu a vzpoměl jsem si na starší přednášku Brandona Rhodese The Clean Architecture in Python (video, slides).

Burying I/O instead of decoupling it

Brandon krásně popisuje jednu běžnou architektonickou chybu, kterou my programátoři děláme. Vezměme si jako příklad úkol: Udělat nějakou akci se všemi soubory s příponou .txt v adresáři. Mohlo by to vypadat třeba takto:

for i in os.listdir(dir):
    if i.endswith(".txt"):
        perform_operation(os.path.join(dir, i))

Ano, ten vstup nebo výstup jsme tady zakopali někam dovnitř aplikace, ikdyž popravdě namítnete, že celý kód je tak krátký, až je přehledný. Co když budeme chtít něco víc? Například jít do podadresářů?

def process_directory(dir):
    for i in os.listdir(dir):
        full_fn = os.path.join(dir, i)
        if os.path.isdir(full_fn):
            process_directory(full_fn)
        elif i.endswith(".txt"):
            perform_operation(full_fn)

Zakopanější a rekurzivní. Standardní knihovna pythonu nám tady nabízí konstrukci os.walk. S její pomocí se zbavíme rekurze (přesněji řečeno os.walk ji udělá za nás). K dokonalosti ale kousek chybí:

for (root, dirs, files) in os.walk(dir):
    for i in files:
        if i.endswith(".txt"):
            perform_operation(os.path.join(dir, root, i))

Pořád tu máme blok, který dělá I/O, a uvnitř v něm je provedení té vlastní funkcionality. Tak ještě jeden pokus:

def select_files_of_interest(dir):
    for (root, dirs, files) in os.walk(dir):
        if i.endswith(".txt"):
            yield os.path.join(dir, root, i)


for full_fn in select_files_of_interest(dir):
    perform_operation(full_fn)

Tady už jsme úplně oddělili to procházení adresářů - získání souborů, které nás zajímají - od vlastního zpracovnání. A totéž teď uděláme pro zip.

A teď dovnitř zipu

Napíšeme funkci deep_walk, která dostane vstupní bod do souborového systému (adresář, soubor, file-like objekt), bude přes ni iterovat, a tím se dostane k jednotlivým souborům, ať už přímo ze souborového systému, nebo ze zip souborů. Vlastní funkce je jednoduchá. Než abychom si ji nějak dopředu rozebírali se na ni prostě podíváme:

from io import BytesIO
import os
import os.path
import sys
import zipfile

# for python 2, instead of from io import BytesIO:
# from six import BytesIO


def deep_walk(f, fn=None):
    """
    Walks through all files within the fn. Visits files in all subdirectories,
    extracts zip files and processes all files with the zip files the same
    way (that means, if there is a zip file inside a zip file, it is extracted
    and processed as well.

    :param f: The starting file or directory or a seekable file-like object
              open for reading.
    :param fn: The file name
    :returns: Yields (filename, file_object) for each file within f. Walks
              into directories, extracts zip files on the way.
    """
    if hasattr(f, "read"):  # a file-like object
        if zipfile.is_zipfile(f):
            zf = zipfile.ZipFile(f, mode="r", allowZip64=True)
            for i in zf.infolist():
                with zf.open(i) as sub_f:
                    for ii in deep_walk(BytesIO(sub_f.read()),
                                        "%s:%s" % (fn, i.filename)):
                        yield ii
        else:  # just return an already opened file
            yield (fn, f)
    elif os.path.isdir(f):
        for (root, dirs, files) in os.walk(f):
            for i in files:
                full_path = os.path.join(root, i)
                with open(full_path, "rb") as sub_f:
                    for ii in deep_walk(sub_f, "%s:%s" % (fn, i)):
                        yield ii
    elif zipfile.is_zipfile(f):
        with open(f, "rb") as zf:
            for ii in deep_walk(zf, f):
                yield ii
    else:
        raise ValueError("The parameter f must be a zipfile or a directory.")

Při použití máme krásně oddělený vstup/výstup od vlastní funkcionality:

for (fn, f_obj) in deep_walk(directory_with_zip_files):
    if fn.endswith(".txt"):
        perform_operation(f_obj)

Poznámky:

• Tohle je generátor.
• Všimněte si, že soubor otevřený ze zipu obalím do BytesIO, než ho rekurzivně zpracovávám. Je to proto, že ZipFile potřebuje seekable stream, neboli objekt, který implementuje funkci seek, což soubor otevřený ze zipu není. Což ale znamená, že se každý takto zpracovávaný soubor načte do paměti!
• Ještě jednou: Všechno běží v paměti, takže počítáme jen se soubory tak malými, že se do paměti vejdou, a to v nejhorším případě i všechny najednou.
• Neřešíme výjimky. Zde záleží na aplikaci, ideální by asi bylo řešení jako má os.walk - argument onerror.
• V souborovém příkladu generátoru jsem vložil ten test na *.txt dovnitř generátoru. Zde nechávám funkci obecnou - vrátí všechny soubory. Můžeme uvážit napsání druhého generátoru, který obalí deep_walk, a bude jen dělat kontrolu na příponu souboru. Ale nahrazovat jednořádkovou podmínku generátorem mi naopak přijde jako neopodstatněné zesložitění.
• Mohli bychom přidat podporu gzipu a dalších formátů.

Slovo o výjimkách

Schválně si zkuste, jak se bude program chovat, když někde vznikne výjimka. Uděláme si takový umělý případ:

def exception_test_generator():
    for i in range(10):
        try:
            yield i
        except Exception:  # catch-all
            print("Caught an exception in the iterator")

for i in exception_test_generator():
    raise Exception()

Tipněte si, jak to dopadne. Zachytí výjimku ten try...except ve funkci exception_test_generator? A je to tak správně? Nebo jej nezachytí? A je to tak správně? Protože nejcennější znalosti jsou ty, které získáte vlastní prací, tak Vám tentokrát odpověď nedám.

Závěr

Když jste začali číst, pravděpodobně Vás napadlo slovo callback. A taky že ano. Callback je další a pravděpodobně starší způsob, jak podobné problémy řešit. Ale zkuste porovnat čitelnost:

# callback:
deep_walk(directory_with_zip_files, callback=process_file)

# generator:
for (fn, f_obj) in deep_walk(directory_with_zip_files):
    process_file(f_obj)

Verze s generátorem rozhodně není kratší, ale pro mě je čitelnější, protože nezakrývá, ale naopak zvýrazňuje tu logiku a strukturu. Kromě toho umožňuje snažší zpracování výjimek. Stejně, jako chceme oddělit vstup a výstup od vlastní funkcionality, chceme od vlastní funkcionality oddělit i zpracování chybových / výjimečných stavů.

Na úplný závěr zopakuju odkaz na přednášku Brandona Rhodese The Clean Architecture in Python (video, slides) a přidám ještě jeden: Raymond Hettinger - Beyond PEP 8 -- Best practices for beautiful intelligible code - PyCon 2015.