<commit_before>import io
import tarfile
import contextlib


class SpoutManager:
    """
    A convenience context manager for creating spouts, allowing you to create
    spout code like:

    ```
    while True:
        with SpoutManager() as spout:
            spout.add_from_bytes("foo", b"#")
        time.sleep(1.0)
    ```
    """

    def __init__(self, marker_filename=None):
        """
        Creates a new spout manager.
        """

        self.f = None
        self.marker_filename = marker_filename

    def __enter__(self):
        self.f = tarfile.open(fileobj=open("/pfs/out", "wb"), mode="w|", encoding="utf-8")
        return self

    def __exit__(self, type, value, traceback):
        self.f.close()

    @contextlib.contextmanager
    def marker(self):
        """
        Gets the marker file as a context manager.
        """

        if self.marker_filename is None:
            raise Exception("no marker filename set")
        with open("/pfs/{}".format(self.marker_filename), "r") as f:
            yield f

    def add_from_fileobj(self, path, size, fileobj):
        """
        Adds a file to the spout from a file-like object.

        Params:

        * `path`: The path to the file in the spout.
        * `size`: The size of the file.
        * `fileobj`: The file-like object to add.
        """

        tar_info = tarfile.TarInfo(path)
        tar_info.size = size
        tar_info.mode = 0o600
        self.f.addfile(tarinfo=tar_info, fileobj=fileobj)

    def add_from_bytes(self, path, bytes):
        """
        Adds a file to the spout from a bytestring.

        Params:

        * `path`: The path to the file in the spout.
        * `bytes`: The bytestring representing the file contents.
        """

        self.add_from_fileobj(path, len(bytes), io.BytesIO(bytes))

    def add_marker_from_fileobj(self, size, fileobj):<commit_msg>Customize output directory in spout managers<commit_after>import io
import os
import tarfile
import contextlib


class SpoutManager:
    """
    A convenience context manager for creating spouts, allowing you to create
    spout code like:

    ```
    while True:
        with SpoutManager() as spout:
            spout.add_from_bytes("foo", b"#")
        time.sleep(1.0)
    ```
    """

    def __init__(self, marker_filename=None, pfs_directory="/pfs"):
        """
        Creates a new spout manager.

        Params:

        * `marker_filename`: The name of the file for storing markers. If
        unspecified, marker-related operations will fail.
        * `pfs_directory`: The directory for PFS content. Usually this
        shouldn't be explicitly specified, unless the spout manager is being
        tested outside of a real Pachyderm pipeline.
        """

        self.f = None
        self.marker_filename = marker_filename
        self.pfs_directory = pfs_directory

    def __enter__(self):
        f = open(os.path.join(self.pfs_directory, "out"), "wb")
        self.f = tarfile.open(fileobj=f, mode="w|", encoding="utf-8")
        return self

    def __exit__(self, type, value, traceback):
        self.f.close()

    @contextlib.contextmanager
    def marker(self):
        """
        Gets the marker file as a context manager.
        """

        if self.marker_filename is None:
            raise Exception("no marker filename set")
        with open(os.path.join(self.pfs_directory, self.marker_filename), "r") as f:
            yield f

    def add_from_fileobj(self, path, size, fileobj):
        """
        Adds a file to the spout from a file-like object.

        Params:

        * `path`: The path to the file in the spout.
        * `size`: The size of the file.
        * `fileobj`: The file-like object to add.
        """

        tar_info = tarfile.TarInfo(path)
        tar_info.size = size
        tar_info.mode = 0o600
        self.f.addfile(tarinfo=tar_info, fileobj=fileobj)

    def add_from_bytes(self, path, bytes):
        """
        Adds a file to the spout from a bytestring.

        Params:

        * `path`: The path to the file in the spout.
        * `bytes`: The bytestring representing the file contents.
        """

        self.add_from_fileobj(path, len(bytes), io.BytesIO(bytes))

    def add_marker_from_fileobj(self, size, fileobj):
{'fraction_non_alphanumeric': 0.10270774976657329, 'fraction_numerical': 0.0032679738562091504, 'mean_word_length': 2.1977611940298507, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4126105', 'n_tokens_mistral': 1321, 'n_tokens_neox': 1261, 'n_words': 388}