<commit_before>import codecs
from collections import namedtuple
import datetime
import tempfile

import boto

from tornado import web

from IPython import nbformat
from IPython.html.services.contents.filecheckpoints import GenericFileCheckpoints
from IPython.html.services.contents.manager import ContentsManager
from IPython.utils import tz


# s3 return different time formats in different situations apparently
S3_TIMEFORMAT_GET_KEY = '%a, %d %b %Y %H:%M:%S GMT'
S3_TIMEFORMAT_BUCKET_LIST = '%Y-%m-%dT%H:%M:%S.000Z'

fakekey = namedtuple('fakekey', 'name')


class S3ContentsManager(ContentsManager):
    @staticmethod
    def _parse_s3_uri(uri, delimiter='/'):
        if not uri.startswith("s3://"):
            raise Exception("Unexpected s3 uri scheme in '{}', expected s3://".format(uri))
        return uri[5:].split(delimiter, 1)

    def _path_to_s3_key(self, path):
        return self.s3_prefix + path.strip(self.s3_key_delimiter)

    def _path_to_s3_key_dir(self, path):
        key = self._path_to_s3_key(path)
        # append delimiter if path is non-empty to avoid s3://bucket//
        if path != '':
            key += self.s3_key_delimiter
        return key

    def _get_key_dir_name(self, name):
        try:
            return name.rsplit(self.s3_key_delimiter, 2)[-2]
        except IndexError:
            return ''

    def _s3_key_dir_to_model(self, key):
        self.log.debug("_s3_key_dir_to_model: %s: %s", key, key.name)
        model = {
            'name': self._get_key_dir_name(key.name),
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.utcnow(), # key.last_modified,  will be used in an HTTP header
            'created': None, # key.last_modified,
            'type': 'directory',
            'content': None,
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_dir_to_model: %s: %s", key.name, model)
        return model

    def _s3_key_file_to_model(self, key, timeformat):
        self.log.debug("_s3_key_file_to_model: %s: %s", key, key.name)
        model = {
            'content': None,
            'name': key.name.rsplit(self.s3_key_delimiter, 1)[-1],
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.strptime(
                key.last_modified, timeformat).replace(tzinfo=tz.UTC),
            'created': None,
            'type': 'file',
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_file_to_model: %s: %s", key.name, model)
        return model

    def _s3_key_notebook_to_model(self, key, timeformat):
        self.log.debug("_s3_key_notebook_to_model: %s: %s", key, key.name)
        model = {
            'content': None,
            'name': key.name.rsplit(self.s3_key_delimiter, 1)[-1],
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.strptime(
                key.last_modified, timeformat).replace(tzinfo=tz.UTC),
            'created': None,
            'type': 'notebook',
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_notebook_to_model: %s: %s", key.name, model)
        return model

    def __init__(self, **kwargs):
        super(S3ContentsManager, self).__init__(**kwargs)
        config = self.config[self.__class__.__name__]  # this still can't be right
        self.s3_base_uri = config['s3_base_uri']
        self.s3_key_delimiter = config.get('s3_key_delimiter', '/')
        self.s3_bucket, self.s3_prefix = self._parse_s3_uri(self.s3_base_uri, self.s3_key_delimiter)
        # ensure prefix ends with the delimiter
        if not self.s3_prefix.endswith(self.s3_key_delimiter) and self.s3_prefix != '':
            self.s3_prefix += self.s3_key_delimiter
        self.s3_connection = boto.connect_s3()
        self.bucket = self.s3_connection.get_bucket(self.s3_bucket)
        self.log.debug("initialized base_uri: %s bucket: %s prefix: %s",
            self.s3_base_uri, self.s3_bucket, self.s3_prefix)

    def _checkpoints_class_default(self):
        return GenericFileCheckpoints

    def list_dirs(self, path):
        self.log.debug('list_dirs: %s', locals())
        key = self._path_to_s3_key_dir(path)
        self.log.debug('list_dirs: looking in bucket:%s under:%s', self.bucket.name, key)
        dirs = []
        for k in self.bucket.list(key, self.s3_key_delimiter):
            if k.name.endswith(self.s3_key_delimiter) and k.name != key:
                dirs.append(self._s3_key_dir_to_model(k))
                self.log.debug('list_dirs: found %s', k.name)
        return dirs

    def list_files(self, path):
        self.log.debug('list_files: %s', locals())
        key = self._path_to_s3_key_dir(path)<commit_msg>Add an option to pass in s3 host<commit_after>import codecs
from collections import namedtuple
import datetime
import tempfile

import boto

from tornado import web

from IPython import nbformat
from IPython.html.services.contents.filecheckpoints import GenericFileCheckpoints
from IPython.html.services.contents.manager import ContentsManager
from IPython.utils import tz


# s3 return different time formats in different situations apparently
S3_TIMEFORMAT_GET_KEY = '%a, %d %b %Y %H:%M:%S GMT'
S3_TIMEFORMAT_BUCKET_LIST = '%Y-%m-%dT%H:%M:%S.000Z'

fakekey = namedtuple('fakekey', 'name')


class S3ContentsManager(ContentsManager):
    @staticmethod
    def _parse_s3_uri(uri, delimiter='/'):
        if not uri.startswith("s3://"):
            raise Exception("Unexpected s3 uri scheme in '{}', expected s3://".format(uri))
        return uri[5:].split(delimiter, 1)

    def _path_to_s3_key(self, path):
        return self.s3_prefix + path.strip(self.s3_key_delimiter)

    def _path_to_s3_key_dir(self, path):
        key = self._path_to_s3_key(path)
        # append delimiter if path is non-empty to avoid s3://bucket//
        if path != '':
            key += self.s3_key_delimiter
        return key

    def _get_key_dir_name(self, name):
        try:
            return name.rsplit(self.s3_key_delimiter, 2)[-2]
        except IndexError:
            return ''

    def _s3_key_dir_to_model(self, key):
        self.log.debug("_s3_key_dir_to_model: %s: %s", key, key.name)
        model = {
            'name': self._get_key_dir_name(key.name),
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.utcnow(), # key.last_modified,  will be used in an HTTP header
            'created': None, # key.last_modified,
            'type': 'directory',
            'content': None,
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_dir_to_model: %s: %s", key.name, model)
        return model

    def _s3_key_file_to_model(self, key, timeformat):
        self.log.debug("_s3_key_file_to_model: %s: %s", key, key.name)
        model = {
            'content': None,
            'name': key.name.rsplit(self.s3_key_delimiter, 1)[-1],
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.strptime(
                key.last_modified, timeformat).replace(tzinfo=tz.UTC),
            'created': None,
            'type': 'file',
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_file_to_model: %s: %s", key.name, model)
        return model

    def _s3_key_notebook_to_model(self, key, timeformat):
        self.log.debug("_s3_key_notebook_to_model: %s: %s", key, key.name)
        model = {
            'content': None,
            'name': key.name.rsplit(self.s3_key_delimiter, 1)[-1],
            'path': key.name.replace(self.s3_prefix, ''),
            'last_modified': datetime.datetime.strptime(
                key.last_modified, timeformat).replace(tzinfo=tz.UTC),
            'created': None,
            'type': 'notebook',
            'mimetype': None,
            'writable': True,
            'format': None,
        }
        self.log.debug("_s3_key_notebook_to_model: %s: %s", key.name, model)
        return model

    def __init__(self, **kwargs):
        super(S3ContentsManager, self).__init__(**kwargs)
        config = self.config[self.__class__.__name__]  # this still can't be right
        self.s3_base_uri = config['s3_base_uri']
        self.s3_host = config['s3_host']
        self.s3_key_delimiter = config.get('s3_key_delimiter', '/')
        self.s3_bucket, self.s3_prefix = self._parse_s3_uri(self.s3_base_uri, self.s3_key_delimiter)
        # ensure prefix ends with the delimiter
        if not self.s3_prefix.endswith(self.s3_key_delimiter) and self.s3_prefix != '':
            self.s3_prefix += self.s3_key_delimiter
        self.s3_connection = boto.connect_s3(host=self.s3_host)
        self.bucket = self.s3_connection.get_bucket(self.s3_bucket)
        self.log.debug("initialized base_uri: %s bucket: %s prefix: %s",
            self.s3_base_uri, self.s3_bucket, self.s3_prefix)

    def _checkpoints_class_default(self):
        return GenericFileCheckpoints

    def list_dirs(self, path):
        self.log.debug('list_dirs: %s', locals())
        key = self._path_to_s3_key_dir(path)
        self.log.debug('list_dirs: looking in bucket:%s under:%s', self.bucket.name, key)
        dirs = []
        for k in self.bucket.list(key, self.s3_key_delimiter):
            if k.name.endswith(self.s3_key_delimiter) and k.name != key:
                dirs.append(self._s3_key_dir_to_model(k))
                self.log.debug('list_dirs: found %s', k.name)
        return dirs

    def list_files(self, path):
        self.log.debug('list_files: %s', locals())
        key = self._path_to_s3_key_dir(path)
{'fraction_non_alphanumeric': 0.12541519879214896, 'fraction_numerical': 0.014091595369904378, 'mean_word_length': 2.5676840215439856, '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': '6147081', 'n_tokens_mistral': 3392, 'n_tokens_neox': 3332, 'n_words': 648}