7
\$\begingroup\$

I implemented a context manager to handle temporary files and directories.

"""Temporary data management."""

from contextlib import ExitStack
from tempfile import template, NamedTemporaryFile, TemporaryDirectory


__all__ = ['TemporaryDataContext']


class TemporaryDataContext(ExitStack):
    """A context manager for temporary files and folders."""

    def __init__(self, root=None, prefix=template):
        """Initializes the context manager with an
        optional temporary data root and prefix.
        """
        super().__init__()
        self.root = root
        self.prefix = prefix

    def mktemp(self, mode='w+b', encoding=None, suffix=''):
        """Creates and registers a named temporary
        file and returns the file object.
        """
        temp = self.enter_context(NamedTemporaryFile(
            suffix=suffix, prefix=self.prefix, dir=self.root))
        temp.file.close()
        return open(temp.name, mode=mode, encoding=encoding)

    def mkdtemp(self, suffix=''):
        """Creates and registers a temporary
        directory and returns its path.
        """
        return self.enter_context(TemporaryDirectory(
            suffix=suffix, prefix=self.prefix, dir=self.root))

It is being invoked by a data processing API to arbitrarily create temporary files and directories, that need to survive the entire processing chain and will be cleaned up at the end.

I am especially interested in optimizing the mktemp method, which seems awkward to me (though it does its job).

\$\endgroup\$
1
  • 1
    \$\begingroup\$ IMO your update would be better as an answer. :) \$\endgroup\$ – Peilonrayz Sep 15 '17 at 8:30
5
\$\begingroup\$

I just realized, that I can simplify mktemp by directly returning the tempfile._TemporaryFileWrapper's buffer, which can be used in a file context for open and closing without deleting the superordinate _TemporaryFileWrapper:

def mktemp(self, mode='w+b', encoding=None, suffix=''):
    """Creates and registers a named temporary
    file and returns the file object.
    """
    return self.enter_context(NamedTemporaryFile(
        mode=mode, encoding=encoding, suffix=suffix, prefix=self.prefix,
        dir=self.root)).file

Update
With the change above, the returned buffer does not behave as expected regarding its attribute name. To fix that, I introduced a wrapper class:

class NamedFile:
    """Wraps a file buffer and its associated name."""

    def __init__(self, file, name):
        """Sets buffer and name."""
        self.file = file
        self.name = name

    def __enter__(self):
        self.file.__enter__()
        return self

    def __exit__(self, typ, value, traceback):
        return self.file.__exit__(typ, value, traceback)

    def __getattr__(self, attr):
        return getattr(self.file, attr)

    @classmethod
    def from_ntf(cls, named_temporary_file):
        """Creates a named file from the specified NamedTemporaryFile."""
        return cls(named_temporary_file.file, named_temporary_file.name)

And return it instead of the raw buffer accordingly:

def mktemp(self, mode='w+b', encoding=None, suffix=''):
    """Creates and registers a named temporary
    file and returns the file object.
    """
    return NamedFile.from_ntf(self.enter_context(NamedTemporaryFile(
        mode=mode, encoding=encoding, suffix=suffix, prefix=self.prefix,
        dir=self.root)))
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.