<commit_before>__all__ = ('BlockStorageInterface',
           'BlockStorageFile'
           'BlockStorageMMapFile',
           'BlockStorageS3')

import os
import mmap
import struct
from multiprocessing.pool import ThreadPool

import boto3
import botocore

from six.moves import xrange

def BlockStorageTypeFactory(storage_type_name):
    if storage_type_name == 'file':
        return BlockStorageFile
    elif storage_type_name == 'mmap':
        return BlockStorageMMapFile
    elif storage_type_name == 's3':
        return BlockStorageS3
    else:
        raise ValueError(
            "BlockStorageTypeFactory: Unsupported storage type: %s"
            % (storage_type_name))

class BlockStorageInterface(object):

    def __enter__(self):
        return self
    def __exit__(self, *args):
        self.close()

    #
    # Abstract Interface
    #

    @classmethod
    def setup(cls, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def user_header_data(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def block_count(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def block_size(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def storage_name(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def close(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def read_blocks(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def read_block(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def write_blocks(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def write_block(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover

class BlockStorageFile(BlockStorageInterface):

    _index_storage_string = "!QQQ"
    _index_offset = struct.calcsize(_index_storage_string)

    def __init__(self, storage_name):
        if not os.path.exists(storage_name):
            raise IOError("Storage location does not exist: %s"
                          % (storage_name))
        self._f = None
        self._storage_name = storage_name
        self._f = open(self.storage_name, "r+b")
        self._f.seek(0)
        self._block_size, self._block_count, self._user_header_size = \
            struct.unpack(
                self._index_storage_string,
                self._f.read(self._index_offset))
        self._user_header_data = bytes()
        if self._user_header_size > 0:
            self._user_header_data = \
                self._f.read(self._user_header_size)
        self._header_offset = self._index_offset + \
                              len(self._user_header_data)

    #
    # Define BlockStorageInterface Methods
    #

    @classmethod
    def setup(cls,
              storage_name=None,
              block_size=None,
              block_count=None,
              initialize=None,
              user_header_data=None,
              ignore_existing=False):
        if (storage_name is None):
            raise ValueError("'storage_name' is required")
        if (block_size is None):
            raise ValueError("'block_size' is required")
        if (block_count is None):
            raise ValueError("'block_count' is required")
        if (not ignore_existing) and \
           os.path.exists(storage_name):
            raise ValueError(
                "Storage location already exists: %s"
                % (storage_name))
        if (block_size <= 0) or (block_size != int(block_size)):
            raise ValueError(
                "Block size (bytes) must be a positive integer: %s"
                % (block_size))
        if (block_count <= 0) or (block_count != int(block_count)):
            raise ValueError(
                "Block count must be a positive integer: %s"
                % (block_count))
        if (user_header_data is not None) and \
           (type(user_header_data) is not bytes):
            raise TypeError(
                "'user_header_data' must be of type bytes. "
                "Invalid type: %s" % (type(user_header_data)))

        if initialize is None:
            zeros = bytes(bytearray(block_size))
            initialize = lambda i: zeros
        with open(storage_name, "wb") as f:
            # create_index
            if user_header_data is None:
                f.write(struct.pack(cls._index_storage_string,
                                    block_size,
                                    block_count,
                                    0))
            else:
                f.write(struct.pack(cls._index_storage_string,
                                    block_size,
                                    block_count,
                                    len(user_header_data)))
                f.write(user_header_data)
            for i in xrange(block_count):
                block = initialize(i)
                assert len(block) == block_size
                f.write(block)
            f.flush()<commit_msg>cleanup the local file if setup fails<commit_after>__all__ = ('BlockStorageInterface',
           'BlockStorageFile'
           'BlockStorageMMapFile',
           'BlockStorageS3')

import os
import mmap
import struct
from multiprocessing.pool import ThreadPool

import boto3
import botocore

from six.moves import xrange

def BlockStorageTypeFactory(storage_type_name):
    if storage_type_name == 'file':
        return BlockStorageFile
    elif storage_type_name == 'mmap':
        return BlockStorageMMapFile
    elif storage_type_name == 's3':
        return BlockStorageS3
    else:
        raise ValueError(
            "BlockStorageTypeFactory: Unsupported storage type: %s"
            % (storage_type_name))

class BlockStorageInterface(object):

    def __enter__(self):
        return self
    def __exit__(self, *args):
        self.close()

    #
    # Abstract Interface
    #

    @classmethod
    def setup(cls, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def user_header_data(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def block_count(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def block_size(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    @property
    def storage_name(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def close(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def read_blocks(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def read_block(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def write_blocks(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover
    def write_block(self, *args, **kwds):
        raise NotImplementedError                      # pragma: no cover

class BlockStorageFile(BlockStorageInterface):

    _index_storage_string = "!QQQ"
    _index_offset = struct.calcsize(_index_storage_string)

    def __init__(self, storage_name):
        if not os.path.exists(storage_name):
            raise IOError("Storage location does not exist: %s"
                          % (storage_name))
        self._f = None
        self._storage_name = storage_name
        self._f = open(self.storage_name, "r+b")
        self._f.seek(0)
        self._block_size, self._block_count, self._user_header_size = \
            struct.unpack(
                self._index_storage_string,
                self._f.read(self._index_offset))
        self._user_header_data = bytes()
        if self._user_header_size > 0:
            self._user_header_data = \
                self._f.read(self._user_header_size)
        self._header_offset = self._index_offset + \
                              len(self._user_header_data)

    #
    # Define BlockStorageInterface Methods
    #

    @classmethod
    def setup(cls,
              storage_name=None,
              block_size=None,
              block_count=None,
              initialize=None,
              user_header_data=None,
              ignore_existing=False):
        if (storage_name is None):
            raise ValueError("'storage_name' is required")
        if (block_size is None):
            raise ValueError("'block_size' is required")
        if (block_count is None):
            raise ValueError("'block_count' is required")
        if (not ignore_existing) and \
           os.path.exists(storage_name):
            raise ValueError(
                "Storage location already exists: %s"
                % (storage_name))
        if (block_size <= 0) or (block_size != int(block_size)):
            raise ValueError(
                "Block size (bytes) must be a positive integer: %s"
                % (block_size))
        if (block_count <= 0) or (block_count != int(block_count)):
            raise ValueError(
                "Block count must be a positive integer: %s"
                % (block_count))
        if (user_header_data is not None) and \
           (type(user_header_data) is not bytes):
            raise TypeError(
                "'user_header_data' must be of type bytes. "
                "Invalid type: %s" % (type(user_header_data)))

        if initialize is None:
            zeros = bytes(bytearray(block_size))
            initialize = lambda i: zeros
        try:
            with open(storage_name, "wb") as f:
                # create_index
                if user_header_data is None:
                    f.write(struct.pack(cls._index_storage_string,
                                        block_size,
                                        block_count,
                                        0))
                else:
                    f.write(struct.pack(cls._index_storage_string,
                                        block_size,
                                        block_count,
                                        len(user_header_data)))
                    f.write(user_header_data)
                for i in xrange(block_count):
                    block = initialize(i)
                    assert len(block) == block_size
                    f.write(block)
                f.flush()
        except:
            os.remove(storage_name)
            raise
{'fraction_non_alphanumeric': 0.08180730652433975, 'fraction_numerical': 0.001656390908254348, 'mean_word_length': 1.6624203821656052, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '948487', 'n_tokens_mistral': 2908, 'n_tokens_neox': 2825, 'n_words': 703}