~didrocks/unity/altf10

__all__ = [

    'MemoryStore',

    ]

import email

import simplejson

from urlparse import parse_qs

from grackle.error import (

    ArchiveIdExists,

    ArchiveIdNotFound,

    MessageIdNotFound,

    UnparsableDateRange,

    UnsupportedDisplayType,

    UnsupportedOrder,

    )

SUPPORTED_DISPLAY_TYPES = set(['all', 'text-only', 'headers-only'])

SUPPORTED_ORDERS = set(

    ['date', 'author', 'subject', 'thread_newest', 'thread_oldest',

     'thread_subject'])

def threaded_messages(messages):

    threads = {}

    count = 0

    pending = []

    for message in messages:

        if message.get('replies') is None:

            threads[message['message_id']] = [message]

            count += 1

        else:

            pending.append(message)

    for message in pending:

        threads[message['replies']].append(message)

    return threads.values()

def get_body_text(message):

    """Return the first plain/text messaage part."""

    if not message.is_multipart():

        return message.get_payload()

    for part in email.iterators.typed_subpart_iterator(message, 'multipart'):

        subparts = part.get_payload()

        for subpart in subparts:

            if subpart.get_content_type() == 'text/plain':

                return subpart.get_payload().strip()

    return ''

def make_json_message(message_id, raw_message, hidden=False):

    message = email.message_from_string(raw_message)

    headers = dict(message.items())

    message = {

        'message_id': message_id,

        'headers': headers,

        # This is broken because the in-reply-to must be encoded.

        # X-Message-ID-Hash is calculated from the Base 32.

        'thread_id': headers.get('in-reply-to', message_id),

        'date': headers.get('date'),

        'subject': headers.get('subject'),

        'author': headers.get('from'),

        'hidden': hidden,

        'attachments': [],

        'replies': headers.get('in-reply-to'),

        'body': get_body_text(message),

        }

    return message

class MemoryStore:

    """A memory-backed message store."""

    def __init__(self, message_archives):

        """Constructor."""

        self.message_archives = message_archives

    @staticmethod

    def is_multipart(message):

        return isinstance(message['body'], list)

    def put_archive(self, archive_id, raw_archive=None):

        # XXX sinzui 2012-02-29: this needs to raise an error

        # if the th archive_id is invalid, or the raw archive is not mbox.

        if archive_id in self.message_archives:

            raise ArchiveIdExists()

        self.message_archives[archive_id] = []

    def put_message(self, archive_id, message_id, raw_message):

        # XXX sinzui 2012-02-29: this needs to raise an error

        # if the th archive_id is invalid, message_id is not base32

        # or the raw message is not an email.

        if archive_id not in self.message_archives:

            raise ArchiveIdNotFound()

        if not raw_message:

            raise ValueError('raw_message is not a message.')

        json_message = make_json_message(message_id, raw_message)

        messages = self.message_archives[archive_id]

        messages.append(json_message)

    def get_messages(self, archive_id, query_string):

        """Return matching messages.

        :param archive_id: The archive to retrieve from.

        :param query_string: Contains 'parameters', which is a JSON-format

            string describing parameters.

        """

        query = parse_qs(query_string)

        parameters = simplejson.loads(query['parameters'][0])

        order = parameters.get('order')

        messages = self.message_archives[archive_id]

        if order is not None:

            if order not in SUPPORTED_ORDERS:

                raise UnsupportedOrder

            elif order.startswith('thread_'):

                threaded = threaded_messages(messages)

                messages = []

                if order == 'thread_subject':

                    threaded.sort(key=lambda t: t[0]['subject'])

                if order == 'thread_oldest':

                    threaded.sort(key=lambda t: min(m['date'] for m in t))

                if order == 'thread_newest':

                    threaded.sort(key=lambda t: max(m['date'] for m in t))

                for thread in threaded:

                    messages.extend(thread)

            else:

                messages.sort(key=lambda m: m[order])

        display_type = parameters.get('display_type', 'all')

        if display_type not in SUPPORTED_DISPLAY_TYPES:

            raise UnsupportedDisplayType

        if 'date_range' in parameters:

            try:

                start_date, end_date = parameters['date_range'].split('..')

                if not start_date or not end_date:

                    raise UnparsableDateRange

            except ValueError:

                raise UnparsableDateRange

        new_messages = []

        for message in messages:

            if (not parameters['include_hidden'] and message['hidden']):

                continue

            if ('message_ids' in parameters

                and message['message_id'] not in parameters['message_ids']):

                continue

            if ('date_range' in parameters

                and (message['date'] < start_date

                     or message['date'] > end_date)):

                continue

            message = dict(message)

            if 'headers' in parameters:

                headers = dict(

                    (k, v) for k, v in message['headers'].iteritems()

                    if k in parameters['headers'])

                message['headers'] = headers

            if display_type == 'headers-only':

                del message['body']

            elif display_type == 'text-only' and self.is_multipart(message):

                text_parts = [

                    part.get_payload() for part in message['body']

                    if part.get_content_type() == 'text/plain']

                message['body'] = '\n\n'.join(text_parts)

            elif display_type == 'all' and self.is_multipart(message):

                parts = [str(part.get_payload()) for part in message['body']]

                message['body'] = '\n\n'.join(parts)

            max_body = parameters.get('max_body_length')

            if max_body is not None and display_type != 'headers-only':

                message['body'] = message['body'][:max_body]

            new_messages.append(message)

        messages = new_messages

        limit = parameters.get('limit', 100)

        memo = parameters.get('memo')

        message_id_indices = dict(

            (m['message_id'], idx) for idx, m in enumerate(messages))

        if memo is None:

            start = 0

        else:

            start = message_id_indices[memo.encode('rot13')]

        if start > 0:

            previous_memo = messages[start - 1]['message_id'].encode('rot13')

        else:

            previous_memo = None

        end = min(start + limit, len(messages))

        if end < len(messages):

            next_memo = messages[end]['message_id'].encode('rot13')

        else:

            next_memo = None

        messages = messages[start:end]

        response = {

            'messages': messages,

            'next_memo': next_memo,

            'previous_memo': previous_memo

            }

        return response

    def hide_message(self, archive_id, message_id, query_string):

        """Change the visbility of a message in an archive.

        :param archive_id: The archive to retrieve from.

        :param query_string: Contains 'parameters', which is a JSON-format

            string describing parameters.

        """

        query = parse_qs(query_string)

        parameters = simplejson.loads(query['parameters'][0])

        hidden = parameters['hidden']

        messages = self.message_archives[archive_id]

        for message in messages:

            if message['message_id'] == message_id:

                message['hidden'] = hidden

            response = {

                'message_id': message_id,

                'hidden': hidden,

                }

            return response

        raise MessageIdNotFound