<commit_before>from __future__ import absolute_import

import boto
import fnmatch
import logging
from io import BytesIO, StringIO

from ...utils import Tokenizer
from .file_system import FileSystem

log = logging.getLogger(__name__)


class S3(FileSystem):
    _conn = None

    def __init__(self, file_name):
        FileSystem.__init__(self, file_name)

        # obtain key
        t = Tokenizer(self.file_name)
        t.next('//')  # skip scheme
        bucket_name = t.next('/')
        key_name = t.next()
        conn = S3._get_conn()
        bucket = conn.get_bucket(bucket_name, validate=False)
        self.key = bucket.get_key(key_name)
        if not self.key:
            self.key = bucket.new_key(key_name)

    @staticmethod
    def _get_conn():
        if not S3._conn:
            S3._conn = boto.connect_s3()
        return S3._conn

    @staticmethod
    def exists(path):
        t = Tokenizer(path)
        t.next('//')  # skip scheme
        bucket_name = t.next('/')
        key_name = t.next()
        conn = S3._get_conn()
        bucket = conn.get_bucket(bucket_name, validate=False)
        return (bucket.get_key(key_name) or
                any(True for _ in bucket.list(prefix=key_name+'/')))

    @staticmethod
    def resolve_filenames(expr):
        files = []

        t = Tokenizer(expr)
        scheme = t.next('://')
        bucket_name = t.next('/')
        prefix = t.next(['*', '?'])

        bucket = S3._get_conn().get_bucket(
            bucket_name,
            validate=False
        )
        expr_after_bucket = expr[len(scheme)+3+len(bucket_name)+1:]
        for k in bucket.list(prefix=prefix):
            if fnmatch.fnmatch(k.name, expr_after_bucket) or \
               fnmatch.fnmatch(k.name, expr_after_bucket+'/part*'):
                files.append(scheme+'://'+bucket_name+'/'+k.name)
        return files

    def load(self):
        log.debug('Loading {0} with size {1}.'
                  ''.format(self.key.name, self.key.size))
        return BytesIO(self.key.get_contents_as_string())

    def load_text(self):
        log.debug('Loading {0} with size {1}.'
                  ''.format(self.key.name, self.key.size))
        return StringIO(self.key.get_contents_as_string())

    def dump(self, stream):
        log.debug('Dumping to {0}.'.format(self.key.name))
        self.key.set_contents_from_string(b''.join(stream))
        return self

    def make_public(self, recursive=False):
        self.key.make_public(recursive)
        return self<commit_msg>bugfix: data type for load_text() on s3 without knowing the encoding is still bytes and not str<commit_after>from __future__ import absolute_import

import boto
import fnmatch
import logging
from io import BytesIO, StringIO

from ...utils import Tokenizer
from .file_system import FileSystem

log = logging.getLogger(__name__)


class S3(FileSystem):
    _conn = None

    def __init__(self, file_name):
        FileSystem.__init__(self, file_name)

        # obtain key
        t = Tokenizer(self.file_name)
        t.next('//')  # skip scheme
        bucket_name = t.next('/')
        key_name = t.next()
        conn = S3._get_conn()
        bucket = conn.get_bucket(bucket_name, validate=False)
        self.key = bucket.get_key(key_name)
        if not self.key:
            self.key = bucket.new_key(key_name)

    @staticmethod
    def _get_conn():
        if not S3._conn:
            S3._conn = boto.connect_s3()
        return S3._conn

    @staticmethod
    def exists(path):
        t = Tokenizer(path)
        t.next('//')  # skip scheme
        bucket_name = t.next('/')
        key_name = t.next()
        conn = S3._get_conn()
        bucket = conn.get_bucket(bucket_name, validate=False)
        return (bucket.get_key(key_name) or
                any(True for _ in bucket.list(prefix=key_name+'/')))

    @staticmethod
    def resolve_filenames(expr):
        files = []

        t = Tokenizer(expr)
        scheme = t.next('://')
        bucket_name = t.next('/')
        prefix = t.next(['*', '?'])

        bucket = S3._get_conn().get_bucket(
            bucket_name,
            validate=False
        )
        expr_after_bucket = expr[len(scheme)+3+len(bucket_name)+1:]
        for k in bucket.list(prefix=prefix):
            if fnmatch.fnmatch(k.name, expr_after_bucket) or \
               fnmatch.fnmatch(k.name, expr_after_bucket+'/part*'):
                files.append(scheme+'://'+bucket_name+'/'+k.name)
        return files

    def load(self):
        log.debug('Loading {0} with size {1}.'
                  ''.format(self.key.name, self.key.size))
        return BytesIO(self.key.get_contents_as_string())

    def load_text(self):
        log.debug('Loading {0} with size {1}.'
                  ''.format(self.key.name, self.key.size))
        return BytesIO(self.key.get_contents_as_string())

    def dump(self, stream):
        log.debug('Dumping to {0}.'.format(self.key.name))
        self.key.set_contents_from_string(b''.join(stream))
        return self

    def make_public(self, recursive=False):
        self.key.make_public(recursive)
        return self
{'fraction_non_alphanumeric': 0.12658474741564268, 'fraction_numerical': 0.006046420908913595, 'mean_word_length': 2.2353312302839115, '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': '6482799', 'n_tokens_mistral': 1648, 'n_tokens_neox': 1611, 'n_words': 333}