Ještě jednou a lépe: práce v dočasně vytvořeném adresáři v Pythonu

Občas se vám může hodit vytvořit dočasný adresář, vstoupit do něj, vykonat určité akce, vrátit se zpět do původního adresáře a dočasný adresář zrušit. V dnešním příspěvku se podíváme na to, jak to (ne)dělat v Pythonu.

Původní kód

import os
import tempfile
 
old_wd = os.getcwd()
with tempfile.TemporaryDirectory() as tmp_dir:
    os.chdir(tmp_dir)
    try:
        # Perform actions in the temporary directory...
    finally:
        os.chdir(old_wd)

Nejdříve si uložíme adresář, ve kterém se nacházíme. Následně využijeme tempfile.TemporaryDirectory k vytvoření dočasného adresáře. V kombinaci s příkazem with dojde jak k jeho vytvoření, tak ke smazání po opuštění bloku. Uvnitř příkazu with se přesuneme do dočasného adresáře, provedeme potřebné akce a adresář opustíme. Blok try...finally je nutný z toho důvodu, abychom se do původního adresáře vrátili i v případě, kdy v bloku # Perform actions... dojde k vyhození výjimky.

Proč takto raději ne?

Ač řešení funguje, tak vede k duplikaci kódu. Pokud byste totiž tuto funkcionalitu potřebovali na více místech, tak na všech byste si museli uložit aktuální adresář, vytvořit dočasný adresář atd. Navíc kód obsahuje logiku pro vytváření adresáře, přesun do něj atd., což jej znepřehledňuje. Chtělo by to nějaké univerzálnější řešení.

Jak to udělat lépe?

Ukáži dvě možná řešení. Obě jsou založena na vytvoření context manageru, který bude poskytovat žádanou funkcionalitu.

První řešení

import os
import tempfile
import shutil
 
class InsideTemporaryDirectory():
    def __enter__(self):
        self.old_wd = os.getcwd()
        self.tmp_dir = tempfile.mkdtemp()
        try:
            os.chdir(self.tmp_dir)
        except:
            shutil.rmtree(self.tmp_dir)
            raise
        return self.tmp_dir
 
    def __exit__(self, *exc_info):
        try:
            os.chdir(self.old_wd)
        finally:
            shutil.rmtree(self.tmp_dir)

Použití:

with InsideTemporaryDirectory() as tmp_dir:
    # Perform actions in the temporary directory...

Vytvořili jsme si context manager svépomoci. V metodě __enter__(), která se vykoná po InsideTemporaryDirectory(), si vytvoříme dočasný adresář s využitím tempfile.mkdtemp, vstoupíme do něj a vrátíme jeho název (cestu). Ta se při použití uloží do tmp_dir. V metodě __exit__() se pak přesuneme zpátky do původního adresáře a původní adresář zrušíme. Onen try...except blok v __enter__() je nutný z toho důvodu, abychom dočasný adresář zrušili, pokud neuspěje přesun do tohoto adresáře, tedy os.chdir(self.tmp_dir). Z podobného důvodu je nutné použít try...finally blok v __exit__().

Druhé řešení

import contextlib
import os
import tempfile
 
@contextlib.contextmanager
def inside_temporary_directory():
    old_wd = os.getcwd()
    with tempfile.TemporaryDirectory() as tmp_dir:
        os.chdir(tmp_dir)
        try:
            yield tmp_dir
        finally:
            os.chdir(old_wd)

Použití:

with inside_temporary_directory() as tmp_dir:
    # Perform actions in the temporary directory...

Místo vytváření context manageru svépomoci využijeme dekorátor contextlib.contextmanager, díky kterému nemusíme definovat třídu s metodami __enter__() a __exit__(), ale stačí nám vytvořit funkci, která vrátí výsledek přes yield. Kód před yield se vykoná při zavolání inside_temporary_directory(), výsledek se uloží do tmp_dir a při ukončení bloku with se vykoná kód za yield. Co je potřeba mít na paměti, tak je, že yield musíme obalit do try...finally bloku, aby došlo k návratu do původního adresáře i v případě, kdy v kódu # Perform actions... dojde k vyhození výjimky. Jako poslední bych zmínil, že za zrušení adresáře odpovídá with tempfile.TemporaryDirectory().

Kód u druhého řešení je díky využití contextlib.contextmanager o něco kratší.

Závěr

V příspěvku jsme viděli, že za cenu jednorázového zvětšení množství kódu lze výsledný kód zpřehlednit a vyhnout se duplikaci kódu. Pokud víte o elegantnějším způsobu, jak docílit toho, čeho jsme se v tomto příspěvku snažili docílit, určitě se ozvěte do komentáře!