<commit_before>""" Helper functions for actions that operate on accounts.

These could be methods of ImapAccount, but separating them gives us more
flexibility with calling code, as most don't need any attributes of the
account object other than the ID, to limit the action.
"""
import os
import json

from hashlib import sha256

from sqlalchemy import distinct, func
from sqlalchemy.orm.exc import NoResultFound

from flanker import mime

from inbox.util.misc import or_none
from inbox.util.addr import parse_email_address
from inbox.util.file import mkdirp
from inbox.util.misc import parse_ml_headers

from .tables import Block, Message, ImapUid, UIDValidity, FolderItem, Thread

from ..config import config
from ..log import get_logger
log = get_logger()

def total_stored_data(account_id, session):
    """ Computes the total size of the block data of emails in your
        account's IMAP folders
    """
    subq = session.query(Block) \
            .join(Block.message, Message.imapuids) \
            .filter(ImapUid.imapaccount_id==account_id) \
            .group_by(Message.id, Block.id)
    return session.query(func.sum(subq.subquery().columns.size)).scalar()

def total_stored_messages(account_id, session):
    """ Computes the number of emails in your account's IMAP folders """
    return session.query(Message) \
            .join(Message.imapuids) \
            .filter(ImapUid.imapaccount_id==account_id) \
            .group_by(Message.id).count()

def all_uids(account_id, session, folder_name):
    return [uid for uid, in
            session.query(ImapUid.msg_uid).filter_by(
                imapaccount_id=account_id, folder_name=folder_name)]

def g_msgids(account_id, session, in_=None):
    query = session.query(distinct(Message.g_msgid)).join(ImapUid) \
                .filter(ImapUid.imapaccount_id==account_id)
    if in_:
        query = query.filter(Message.g_msgid.in_(in_))
    return sorted([g_msgid for g_msgid, in query], key=long)

def g_metadata(account_id, session, folder_name):
    query = session.query(ImapUid.msg_uid, Message.g_msgid,
                Message.g_thrid).filter(
                        ImapUid.imapaccount_id==account_id,
                        ImapUid.folder_name==folder_name,
                        ImapUid.message_id==Message.id)

    return dict([(int(uid), dict(msgid=g_msgid, thrid=g_thrid)) \
            for uid, g_msgid, g_thrid in query])

def update_metadata(account_id, session, folder_name, uids, new_flags):
    """ Update flags (the only metadata that can change). """
    for item in session.query(ImapUid).filter(
            ImapUid.imapaccount_id==account_id,
            ImapUid.msg_uid.in_(uids),
            ImapUid.folder_name==folder_name):
        flags = new_flags[item.msg_uid]['flags']
        labels = new_flags[item.msg_uid]['labels']
        item.update_flags(flags, labels)

def remove_messages(account_id, session, uids, folder):
    fm_query = session.query(ImapUid).filter(
            ImapUid.imapaccount_id==account_id,
            ImapUid.folder_name==folder,
            ImapUid.msg_uid.in_(uids))
    fm_query.delete(synchronize_session='fetch')

    # XXX TODO: Have a recurring worker permanently remove dangling
    # messages from the database and block store. (Probably too
    # expensive to do here.)

def get_uidvalidity(account_id, session, folder_name):
    try:
        # using .one() here may catch duplication bugs
        return session.query(UIDValidity).filter_by(
                imapaccount_id=account_id, folder_name=folder_name).one()
    except NoResultFound:
        return None

def uidvalidity_valid(account_id, session, selected_uidvalidity, \
        folder_name, cached_uidvalidity=None):
    """ Validate UIDVALIDITY on currently selected folder. """
    if cached_uidvalidity is None:
        cached_uidvalidity = get_uidvalidity(account_id,
                session, folder_name).uid_validity
        assert type(cached_uidvalidity) == type(selected_uidvalidity), \
                "cached_validity: {0} / selected_uidvalidity: {1}".format(
                        type(cached_uidvalidity),
                        type(selected_uidvalidity))

    if cached_uidvalidity is None:
        # no row is basically equivalent to UIDVALIDITY == -inf
        return True
    else:
        return selected_uidvalidity >= cached_uidvalidity

def update_uidvalidity(account_id, session, folder_name, uidvalidity,
        highestmodseq):
    cached_validity = get_uidvalidity(account_id, session, folder_name)
    if cached_validity is None:
        cached_validity = UIDValidity(imapaccount_id=account_id,
                folder_name=folder_name)
    cached_validity.highestmodseq = highestmodseq
    cached_validity.uid_validity = uidvalidity
    session.add(cached_validity)

def create_message(db_session, log, account, folder_name, uid, internaldate,
        flags, body):
    """ Parses message data, creates metadata database entries, and writes mail
        parts to disk.

        Returns the new ImapUid, which links to new Message and Block
        objects through relationships. All new objects are uncommitted.

        Threads are not computed here; you gotta do that separately.

        This is the one function in this file that gets to take an account
        object instead of an account_id, because we need to relate the
        account to ImapUids for versioning to work, since it needs to look
        up the namespace.
    """
    # trickle-down bugs
    assert account is not None and account.namespace is not None
    try:
        parsed = mime.from_string(body)

        mime_version = parsed.headers.get('Mime-Version')
        # NOTE: sometimes MIME-Version is set to "1.0 (1.0)", hence the .startswith
        if mime_version is not None and not mime_version.startswith('1.0'):
            log.error("Unexpected MIME-Version: %s" % mime_version)

        new_msg = Message()
        new_msg.data_sha256 = sha256(body).hexdigest()

        # clean_subject strips re:, fwd: etc.
        new_msg.subject = parsed.clean_subject
        new_msg.from_addr = parse_email_address(parsed.headers.get('From'))
        new_msg.sender_addr = parse_email_address(parsed.headers.get('Sender'))
        new_msg.reply_to = parse_email_address(parsed.headers.get('Reply-To'))
        new_msg.to_addr = or_none(parsed.headers.getall('To'),
                lambda tos: filter(lambda p: p is not None,
                    [parse_email_address(t) for t in tos]))
        new_msg.cc_addr = or_none(parsed.headers.getall('Cc'),
                lambda ccs: filter(lambda p: p is not None,
                    [parse_email_address(c) for c in ccs]))
        new_msg.bcc_addr = or_none(parsed.headers.getall('Bcc'),
                lambda bccs: filter(lambda p: p is not None,
                    [parse_email_address(c) for c in bccs]))
        new_msg.in_reply_to = parsed.headers.get('In-Reply-To')
        new_msg.message_id = parsed.headers.get('Message-Id')

        new_msg.internaldate = internaldate

        # Optional mailing list headers
        new_msg.mailing_list_headers = parse_ml_headers(parsed.headers)

        imapuid = ImapUid(imapaccount=account, folder_name=folder_name,
                msg_uid=uid, message=new_msg)
        imapuid.update_flags(flags)

        new_msg.size = len(body)  # includes headers text

        i = 0  # for walk_index

        # Store all message headers as object with index 0
        headers_part = Block()
        headers_part.message = new_msg
        headers_part.walk_index = i
        headers_part._data = json.dumps(parsed.headers.items())
        headers_part.size = len(headers_part._data)
        headers_part.data_sha256 = sha256(headers_part._data).hexdigest()
        new_msg.parts.append(headers_part)

        for mimepart in parsed.walk(
                with_self=parsed.content_type.is_singlepart()):
            err = open('err.txt', 'w')
            err.write(mimepart.to_string())
            err.close()
            i += 1
            if mimepart.content_type.is_multipart():
                log.warning("multipart sub-part found! on {0}".format(new_msg.g_msgid))
                continue  # TODO should we store relations?

            new_part = Block()
            new_part.message = new_msg
            new_part.walk_index = i
            new_part.misc_keyval = mimepart.headers.items()  # everything
            new_part.content_type = mimepart.content_type.value
            new_part.filename = mimepart.content_type.params.get('name')

            # Content-Disposition attachment; filename="floorplan.gif"
            if mimepart.content_disposition[0] is not None:
                value, params = mimepart.content_disposition
                if value not in ['inline', 'attachment']:
                    errmsg = """
    Unknown Content-Disposition on message {0} found in {1}.
    Bad Content-Disposition was: '{2}'
    Parsed Content-Disposition was: '{3}'""".format(uid, folder_name,
        mimepart.content_disposition)
                    log.error(errmsg)
                    continue
                else:
                    new_part.content_disposition = value
                    if value == 'attachment':
                        new_part.filename = params.get('filename')

            if mimepart.body is None:
                data_to_write = ''
            elif new_part.content_type.startswith('text'):
                data_to_write = mimepart.body.encode('utf-8', 'strict')
            else:
                data_to_write = mimepart.body
            if data_to_write is None:
                data_to_write = ''
            # normalize mac/win/unix newlines
            data_to_write = data_to_write \
                    .replace('\r\n', '\n').replace('\r', '\n')

            new_part.content_id = mimepart.headers.get('Content-Id')

            new_part._data = data_to_write
            new_part.size = len(data_to_write)
            new_part.data_sha256 = sha256(data_to_write).hexdigest()
            new_msg.parts.append(new_part)
    except mime.DecodingError:
        log_decode_error(account.id, folder_name, uid, body)
        log.error("DecodeError encountered, unparseable message logged to {0}" \
                .format(get_errfilename(account.id, folder_name, uid)))
        return
    new_msg.calculate_sanitized_body()

    return imapuid

def get_errfilename(account_id, folder_name, uid):
    errdir = os.path.join(config['LOGDIR'], str(account_id), 'errors',
            folder_name)
    errfile = os.path.join(errdir, str(uid))
    mkdirp(errdir)
    return errfile

def log_decode_error(account_id, folder_name, uid, msg_string):
    """ msg_string is in the original encoding pulled off the wire """
    errfile = get_errfilename(account_id, folder_name, uid)
    with open(errfile, 'w') as fh:
        fh.write(msg_string)

def add_gmail_attrs(db_session, log, new_uid, flags, folder_name, x_gm_thrid,
        x_gm_msgid, x_gm_labels):
    """ Gmail-specific post-create-message bits."""

    new_uid.message.g_msgid = x_gm_msgid
    # NOTE: g_thrid == g_msgid on the first message in the thread :)
    new_uid.message.g_thrid = x_gm_thrid
    new_uid.update_flags(flags, x_gm_labels)

    # NOTE: This code _requires_ autoflush=True, otherwise duplicate
    # threads may attempt to be created and crash.
    thread = new_uid.message.thread = Thread.from_message(db_session,
            new_uid.imapaccount.namespace, new_uid.message)
    # make sure this thread has all the correct labels
    existing_labels = set([l.folder_name.lower() for l in thread.folders])
    # convert things like \Inbox -> Inbox, \Important -> Important
    # also, gmail labels are case-insensitive
    new_labels = set([l.lstrip('\\').lower() for l in x_gm_labels] + \
            [folder_name.lower()])
    # remove labels that have been deleted -- note that the \Sent label is
    # per-message, not per-thread, but since we always work at the thread
    # level, _we_ apply the label to the whole thread. same goes for
    # \Important.
    thread.folders = [l for l in thread.folders if l.folder_name in new_labels \
            or l.folder_name in ('sent', 'important')]
    # add new labels
    for label in new_labels:
        if label not in existing_labels:
            item = FolderItem(thread=thread, folder_name=label)
            db_session.add(item)

    return new_uid

def create_gmail_message(db_session, log, account, folder_name, uid,
        internaldate, flags, body, x_gm_thrid, x_gm_msgid, x_gm_labels):
    new_uid = create_message(db_session, log, account, folder_name, uid,
            internaldate, flags, body)
    if new_uid:
        return add_gmail_attrs(db_session, log, new_uid, flags, folder_name,
                x_gm_thrid, x_gm_msgid, x_gm_labels)
<commit_msg>Canonicalize All Mail to be 'archive' in our datastore<commit_after>""" Helper functions for actions that operate on accounts.

These could be methods of ImapAccount, but separating them gives us more
flexibility with calling code, as most don't need any attributes of the
account object other than the ID, to limit the action.
"""
import os
import json

from hashlib import sha256

from sqlalchemy import distinct, func
from sqlalchemy.orm.exc import NoResultFound

from flanker import mime

from inbox.util.misc import or_none
from inbox.util.addr import parse_email_address
from inbox.util.file import mkdirp
from inbox.util.misc import parse_ml_headers

from .tables import Block, Message, ImapUid, UIDValidity, FolderItem, Thread

from ..config import config
from ..log import get_logger
log = get_logger()

def total_stored_data(account_id, session):
    """ Computes the total size of the block data of emails in your
        account's IMAP folders
    """
    subq = session.query(Block) \
            .join(Block.message, Message.imapuids) \
            .filter(ImapUid.imapaccount_id==account_id) \
            .group_by(Message.id, Block.id)
    return session.query(func.sum(subq.subquery().columns.size)).scalar()

def total_stored_messages(account_id, session):
    """ Computes the number of emails in your account's IMAP folders """
    return session.query(Message) \
            .join(Message.imapuids) \
            .filter(ImapUid.imapaccount_id==account_id) \
            .group_by(Message.id).count()

def all_uids(account_id, session, folder_name):
    return [uid for uid, in
            session.query(ImapUid.msg_uid).filter_by(
                imapaccount_id=account_id, folder_name=folder_name)]

def g_msgids(account_id, session, in_=None):
    query = session.query(distinct(Message.g_msgid)).join(ImapUid) \
                .filter(ImapUid.imapaccount_id==account_id)
    if in_:
        query = query.filter(Message.g_msgid.in_(in_))
    return sorted([g_msgid for g_msgid, in query], key=long)

def g_metadata(account_id, session, folder_name):
    query = session.query(ImapUid.msg_uid, Message.g_msgid,
                Message.g_thrid).filter(
                        ImapUid.imapaccount_id==account_id,
                        ImapUid.folder_name==folder_name,
                        ImapUid.message_id==Message.id)

    return dict([(int(uid), dict(msgid=g_msgid, thrid=g_thrid)) \
            for uid, g_msgid, g_thrid in query])

def update_metadata(account_id, session, folder_name, uids, new_flags):
    """ Update flags (the only metadata that can change). """
    for item in session.query(ImapUid).filter(
            ImapUid.imapaccount_id==account_id,
            ImapUid.msg_uid.in_(uids),
            ImapUid.folder_name==folder_name):
        flags = new_flags[item.msg_uid]['flags']
        labels = new_flags[item.msg_uid]['labels']
        item.update_flags(flags, labels)

def remove_messages(account_id, session, uids, folder):
    fm_query = session.query(ImapUid).filter(
            ImapUid.imapaccount_id==account_id,
            ImapUid.folder_name==folder,
            ImapUid.msg_uid.in_(uids))
    fm_query.delete(synchronize_session='fetch')

    # XXX TODO: Have a recurring worker permanently remove dangling
    # messages from the database and block store. (Probably too
    # expensive to do here.)

def get_uidvalidity(account_id, session, folder_name):
    try:
        # using .one() here may catch duplication bugs
        return session.query(UIDValidity).filter_by(
                imapaccount_id=account_id, folder_name=folder_name).one()
    except NoResultFound:
        return None

def uidvalidity_valid(account_id, session, selected_uidvalidity, \
        folder_name, cached_uidvalidity=None):
    """ Validate UIDVALIDITY on currently selected folder. """
    if cached_uidvalidity is None:
        cached_uidvalidity = get_uidvalidity(account_id,
                session, folder_name).uid_validity
        assert type(cached_uidvalidity) == type(selected_uidvalidity), \
                "cached_validity: {0} / selected_uidvalidity: {1}".format(
                        type(cached_uidvalidity),
                        type(selected_uidvalidity))

    if cached_uidvalidity is None:
        # no row is basically equivalent to UIDVALIDITY == -inf
        return True
    else:
        return selected_uidvalidity >= cached_uidvalidity

def update_uidvalidity(account_id, session, folder_name, uidvalidity,
        highestmodseq):
    cached_validity = get_uidvalidity(account_id, session, folder_name)
    if cached_validity is None:
        cached_validity = UIDValidity(imapaccount_id=account_id,
                folder_name=folder_name)
    cached_validity.highestmodseq = highestmodseq
    cached_validity.uid_validity = uidvalidity
    session.add(cached_validity)

def create_message(db_session, log, account, folder_name, uid, internaldate,
        flags, body):
    """ Parses message data, creates metadata database entries, and writes mail
        parts to disk.

        Returns the new ImapUid, which links to new Message and Block
        objects through relationships. All new objects are uncommitted.

        Threads are not computed here; you gotta do that separately.

        This is the one function in this file that gets to take an account
        object instead of an account_id, because we need to relate the
        account to ImapUids for versioning to work, since it needs to look
        up the namespace.
    """
    # trickle-down bugs
    assert account is not None and account.namespace is not None
    try:
        parsed = mime.from_string(body)

        mime_version = parsed.headers.get('Mime-Version')
        # NOTE: sometimes MIME-Version is set to "1.0 (1.0)", hence the .startswith
        if mime_version is not None and not mime_version.startswith('1.0'):
            log.error("Unexpected MIME-Version: %s" % mime_version)

        new_msg = Message()
        new_msg.data_sha256 = sha256(body).hexdigest()

        # clean_subject strips re:, fwd: etc.
        new_msg.subject = parsed.clean_subject
        new_msg.from_addr = parse_email_address(parsed.headers.get('From'))
        new_msg.sender_addr = parse_email_address(parsed.headers.get('Sender'))
        new_msg.reply_to = parse_email_address(parsed.headers.get('Reply-To'))
        new_msg.to_addr = or_none(parsed.headers.getall('To'),
                lambda tos: filter(lambda p: p is not None,
                    [parse_email_address(t) for t in tos]))
        new_msg.cc_addr = or_none(parsed.headers.getall('Cc'),
                lambda ccs: filter(lambda p: p is not None,
                    [parse_email_address(c) for c in ccs]))
        new_msg.bcc_addr = or_none(parsed.headers.getall('Bcc'),
                lambda bccs: filter(lambda p: p is not None,
                    [parse_email_address(c) for c in bccs]))
        new_msg.in_reply_to = parsed.headers.get('In-Reply-To')
        new_msg.message_id = parsed.headers.get('Message-Id')

        new_msg.internaldate = internaldate

        # Optional mailing list headers
        new_msg.mailing_list_headers = parse_ml_headers(parsed.headers)

        imapuid = ImapUid(imapaccount=account, folder_name=folder_name,
                msg_uid=uid, message=new_msg)
        imapuid.update_flags(flags)

        new_msg.size = len(body)  # includes headers text

        i = 0  # for walk_index

        # Store all message headers as object with index 0
        headers_part = Block()
        headers_part.message = new_msg
        headers_part.walk_index = i
        headers_part._data = json.dumps(parsed.headers.items())
        headers_part.size = len(headers_part._data)
        headers_part.data_sha256 = sha256(headers_part._data).hexdigest()
        new_msg.parts.append(headers_part)

        for mimepart in parsed.walk(
                with_self=parsed.content_type.is_singlepart()):
            err = open('err.txt', 'w')
            err.write(mimepart.to_string())
            err.close()
            i += 1
            if mimepart.content_type.is_multipart():
                log.warning("multipart sub-part found! on {0}".format(new_msg.g_msgid))
                continue  # TODO should we store relations?

            new_part = Block()
            new_part.message = new_msg
            new_part.walk_index = i
            new_part.misc_keyval = mimepart.headers.items()  # everything
            new_part.content_type = mimepart.content_type.value
            new_part.filename = mimepart.content_type.params.get('name')

            # Content-Disposition attachment; filename="floorplan.gif"
            if mimepart.content_disposition[0] is not None:
                value, params = mimepart.content_disposition
                if value not in ['inline', 'attachment']:
                    errmsg = """
    Unknown Content-Disposition on message {0} found in {1}.
    Bad Content-Disposition was: '{2}'
    Parsed Content-Disposition was: '{3}'""".format(uid, folder_name,
        mimepart.content_disposition)
                    log.error(errmsg)
                    continue
                else:
                    new_part.content_disposition = value
                    if value == 'attachment':
                        new_part.filename = params.get('filename')

            if mimepart.body is None:
                data_to_write = ''
            elif new_part.content_type.startswith('text'):
                data_to_write = mimepart.body.encode('utf-8', 'strict')
            else:
                data_to_write = mimepart.body
            if data_to_write is None:
                data_to_write = ''
            # normalize mac/win/unix newlines
            data_to_write = data_to_write \
                    .replace('\r\n', '\n').replace('\r', '\n')

            new_part.content_id = mimepart.headers.get('Content-Id')

            new_part._data = data_to_write
            new_part.size = len(data_to_write)
            new_part.data_sha256 = sha256(data_to_write).hexdigest()
            new_msg.parts.append(new_part)
    except mime.DecodingError:
        log_decode_error(account.id, folder_name, uid, body)
        log.error("DecodeError encountered, unparseable message logged to {0}" \
                .format(get_errfilename(account.id, folder_name, uid)))
        return
    new_msg.calculate_sanitized_body()

    return imapuid

def get_errfilename(account_id, folder_name, uid):
    errdir = os.path.join(config['LOGDIR'], str(account_id), 'errors',
            folder_name)
    errfile = os.path.join(errdir, str(uid))
    mkdirp(errdir)
    return errfile

def log_decode_error(account_id, folder_name, uid, msg_string):
    """ msg_string is in the original encoding pulled off the wire """
    errfile = get_errfilename(account_id, folder_name, uid)
    with open(errfile, 'w') as fh:
        fh.write(msg_string)

def add_gmail_attrs(db_session, log, new_uid, flags, folder_name, x_gm_thrid,
        x_gm_msgid, x_gm_labels):
    """ Gmail-specific post-create-message bits."""

    new_uid.message.g_msgid = x_gm_msgid
    # NOTE: g_thrid == g_msgid on the first message in the thread :)
    new_uid.message.g_thrid = x_gm_thrid
    new_uid.update_flags(flags, x_gm_labels)

    # NOTE: This code _requires_ autoflush=True, otherwise duplicate
    # threads may attempt to be created and crash.
    thread = new_uid.message.thread = Thread.from_message(db_session,
            new_uid.imapaccount.namespace, new_uid.message)
    # make sure this thread has all the correct labels
    existing_labels = set([l.folder_name.lower() for l in thread.folders])
    # convert things like \Inbox -> Inbox, \Important -> Important
    # also, gmail labels are case-insensitive
    new_labels = set([l.lstrip('\\').lower() for l in x_gm_labels] + \
            [folder_name.lower()])
    # remove labels that have been deleted -- note that the \Sent label is
    # per-message, not per-thread, but since we always work at the thread
    # level, _we_ apply the label to the whole thread. same goes for
    # \Important.
    thread.folders = [l for l in thread.folders if l.folder_name in new_labels \
            or l.folder_name in ('sent', 'important')]
    # canonicalize "All Mail" to be "archive" in Inbox's data representation
    archive_folder_name = new_uid.imapaccount.archive_folder_name.lower()
    if archive_folder_name in new_labels:
        new_labels.remove(archive_folder_name)
        new_labels.add('archive')
    # add new labels
    for label in new_labels:
        if label not in existing_labels:
            item = FolderItem(thread=thread, folder_name=label)
            db_session.add(item)

    return new_uid

def create_gmail_message(db_session, log, account, folder_name, uid,
        internaldate, flags, body, x_gm_thrid, x_gm_msgid, x_gm_labels):
    new_uid = create_message(db_session, log, account, folder_name, uid,
            internaldate, flags, body)
    if new_uid:
        return add_gmail_attrs(db_session, log, new_uid, flags, folder_name,
                x_gm_thrid, x_gm_msgid, x_gm_labels)
<|endoftext|>


<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import collections

from pymatgen.util.testing import PymatgenTest
from pymatgen.core.units import (Energy, Time, Length, unitized, Mass, Memory,
                                 EnergyArray, TimeArray, LengthArray, Unit,
                                 FloatWithUnit, ArrayWithUnit, UnitError)

class UnitTest(PymatgenTest):

    def test_init(self):
        u1 = Unit((("m", 1), ("s", -1)))
        self.assertEqual(str(u1), "m s^-1")
        u2 = Unit("kg m ^ 2 s ^ -2")
        self.assertEqual(str(u2), "J")
        self.assertEqual(str(u1 * u2), "J m s^-1")
        self.assertEqual(str(u2 / u1), "J s m^-1")
        self.assertEqual(str(u1 / Unit("m")), "s^-1")
        self.assertEqual(str(u1 * Unit("s")), "m")

        acc = u1 / Unit("s")
        newton = Unit("kg") * acc
        self.assertEqual(str(newton * Unit("m")), "N m")

class FloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        a = Energy(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = Energy(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        self.assertRaises(UnitError, Energy, 1, "m")

        d = Energy(1, "Ha")
        self.assertAlmostEqual(a + d, 28.311385050630882)
        self.assertAlmostEqual(a - d, -26.11138505063088)
        self.assertEqual(a + 1, 2.1)
        self.assertEqual(str(a / d), "1.1 eV Ha^-1")

    def test_time(self):
        a = Time(20, "h")
        self.assertAlmostEqual(float(a.to("s")), 3600 * 20)
        #Test left and right multiplication.
        b = a * 3
        self.assertAlmostEqual(float(b), 60.0)
        self.assertEqual(str(b.unit), "h")
        self.assertEqual(float(3 * a), 60.0)

    def test_length(self):
        x = Length(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")
        y = x ** 3
        self.assertAlmostEqual(y, 74.088)
        self.assertEqual(str(y.unit), "ang^3")

    def test_memory(self):
        mega = Memory(1, "Mb")
        print(type(Memory), type(mega))
        self.assertEqual(mega.to("byte"), 1024**2)
        self.assertEqual(mega, Memory(1, "mb"))

        same_mega = Memory.from_string("1Mb")
        self.assertEqual(same_mega.unit_type, "memory")

        other_mega = Memory.from_string("+1.0 mb")
        self.assertEqual(mega, other_mega)

    def test_unitized(self):

        @unitized("eV")
        def f():
            return [1, 2, 3]

        self.assertEqual(str(f()[0]), "1.0 eV")
        self.assertIsInstance(f(), list)

        @unitized("eV")
        def g():
            return 2, 3, 4

        self.assertEqual(str(g()[0]), "2.0 eV")
        self.assertIsInstance(g(), tuple)

        @unitized("pm")
        def h():
            d = collections.OrderedDict()
            for i in range(3):
                d[i] = i * 20
            return d

        self.assertEqual(str(h()[1]), "20.0 pm")
        self.assertIsInstance(h(), collections.OrderedDict)

    def test_compound_operations(self):
        g = 10 * Length(1, "m") / (Time(1, "s") ** 2)
        e = Mass(1, "kg") * g * Length(1, "m")
        self.assertEqual(str(e), "10.0 N m")
        form_e = FloatWithUnit(10, unit="kJ mol^-1").to("eV atom^-1")
        self.assertAlmostEqual(float(form_e), 0.103642691905)
        self.assertEqual(str(form_e.unit), "eV atom^-1")
        self.assertRaises(UnitError, form_e.to, "m s^-1")
        a = FloatWithUnit(1.0, "Ha^3")
        b = a.to("J^3")
        self.assertAlmostEqual(b, 8.28672661615e-53)
        self.assertEqual(str(b.unit), "J^3")
        a = FloatWithUnit(1.0, "Ha bohr^-2")
        b = a.to("J m^-2")
        self.assertAlmostEqual(b, 1556.8929672799854)
        self.assertEqual(str(b.unit), "J m^-2")

    def test_as_base_units(self):
        x = FloatWithUnit(5, "MPa")
        self.assert_equal(FloatWithUnit(5000000, "Pa"), x.as_base_units)

class ArrayWithFloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        """
        Similar to FloatWithUnitTest.test_energy.
        Check whether EnergyArray and FloatWithUnit have same behavior.

        # TODO
        One can merge the two tests easily:

        for obj in [Energy, EnergyArray]:
            a = obj(...)
            self.assert(...)

        """
        a = EnergyArray(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(float(b), 0.0404242579378)
        c = EnergyArray(3.14, "J")
        self.assertAlmostEqual(float(c.to("eV")), 1.9598339337836966e+19)
        # self.assertRaises(ValueError, Energy, 1, "m")

        d = EnergyArray(1, "Ha")
        self.assertAlmostEqual(float(a + d), 28.311385050630882)
        self.assertAlmostEqual(float(a - d), -26.11138505063088)
        self.assertEqual(float(a + 1), 2.1)

    def test_time(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        # here there's a minor difference because we have a ndarray with
        # dtype=np.int.
        a = TimeArray(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60 h")
        self.assertEqual(str(3 * a), "60 h")

    def test_length(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        x = LengthArray(4.2, "ang")
        self.assertAlmostEqual(float(x.to("cm")), 4.2e-08)
        self.assertEqual(float(x.to("pm")), 420)
        self.assertEqual(str(x / 2), "2.1 ang")

    def test_array_algebra(self):
        ene_ha = EnergyArray([1, 2], "Ha")
        ene_ev = EnergyArray([1, 2], "eV")
        time_s = TimeArray([1, 2], "s")

        e1 = ene_ha.copy()
        e1 += 1
        e2 = ene_ha.copy()
        e2 -= 1
        e3 = ene_ha.copy()
        #e3 /= 2
        e4 = ene_ha.copy()
        e4 *= 2

        objects_with_unit = [
            ene_ha + ene_ev,
            ene_ha - ene_ev,
            3 * ene_ha,
            ene_ha * 3,
            ene_ha / 3,
            3 / ene_ha,
            ene_ha * time_s,
            ene_ha / ene_ev,
            ene_ha.copy(),
            ene_ha[0:1],
            e1,
            e2,
            e3,
            e4,
        ]

        for i, obj in enumerate(objects_with_unit):
            #print(i, obj.unit)
            self.assertTrue(hasattr(obj, "unit"))
            #self.assertTrue(str(obj.unit) == "Ha")

        objects_without_unit = [
            # Here we could return a FloatWithUnit object but I prefer this
            # a bare scalar since FloatWithUnit extends float while we could
            # have an int.
            ene_ha[0],
        ]

        for obj in objects_without_unit:
            self.assertFalse(hasattr(obj, "unit"))

        with self.assertRaises(UnitError):
            ene_ha + time_s

    def test_factors(self):
        e = EnergyArray([27.21138386, 1], "eV").to("Ha")
        self.assertTrue(str(e) == "[ 0.99999996  0.03674932] Ha")
        l = LengthArray([1.0], "ang").to("bohr")
        self.assertTrue(str(l) == "[ 1.88972612] bohr")
        v = ArrayWithUnit([1, 2, 3], "bohr^3").to("ang^3")
        self.assertTrue(str(v) == '[ 0.14818471  0.29636942  0.44455413] ang^3')

    def test_as_base_units(self):
        x = ArrayWithUnit([5, 10], "MPa")
        self.assert_equal(ArrayWithUnit([5000000, 10000000], "Pa"), x.as_base_units)

class DataPersistenceTest(PymatgenTest):
    def test_pickle(self):
        """Test whether FloatWithUnit and ArrayWithUnit support pickle"""

        for cls in [FloatWithUnit, ArrayWithUnit]:
            a = cls(1, "eV")
            b = cls(10, "N bohr")
            objects = [a, b]

            new_objects_from_protocol = self.serialize_with_pickle(objects)

            for new_objects in new_objects_from_protocol:
                for old_item, new_item in zip(objects, new_objects):
                    self.assertTrue(str(old_item) == str(new_item))

if __name__ == '__main__':
    import unittest
    unittest.main()
<commit_msg>add unit tests for @unitize mods<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import collections

from pymatgen.util.testing import PymatgenTest
from pymatgen.core.units import (Energy, Time, Length, unitized, Mass, Memory,
                                 EnergyArray, TimeArray, LengthArray, Unit,
                                 FloatWithUnit, ArrayWithUnit, UnitError)

class UnitTest(PymatgenTest):

    def test_init(self):
        u1 = Unit((("m", 1), ("s", -1)))
        self.assertEqual(str(u1), "m s^-1")
        u2 = Unit("kg m ^ 2 s ^ -2")
        self.assertEqual(str(u2), "J")
        self.assertEqual(str(u1 * u2), "J m s^-1")
        self.assertEqual(str(u2 / u1), "J s m^-1")
        self.assertEqual(str(u1 / Unit("m")), "s^-1")
        self.assertEqual(str(u1 * Unit("s")), "m")

        acc = u1 / Unit("s")
        newton = Unit("kg") * acc
        self.assertEqual(str(newton * Unit("m")), "N m")

class FloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        a = Energy(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = Energy(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        self.assertRaises(UnitError, Energy, 1, "m")

        d = Energy(1, "Ha")
        self.assertAlmostEqual(a + d, 28.311385050630882)
        self.assertAlmostEqual(a - d, -26.11138505063088)
        self.assertEqual(a + 1, 2.1)
        self.assertEqual(str(a / d), "1.1 eV Ha^-1")

    def test_time(self):
        a = Time(20, "h")
        self.assertAlmostEqual(float(a.to("s")), 3600 * 20)
        #Test left and right multiplication.
        b = a * 3
        self.assertAlmostEqual(float(b), 60.0)
        self.assertEqual(str(b.unit), "h")
        self.assertEqual(float(3 * a), 60.0)

    def test_length(self):
        x = Length(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")
        y = x ** 3
        self.assertAlmostEqual(y, 74.088)
        self.assertEqual(str(y.unit), "ang^3")

    def test_memory(self):
        mega = Memory(1, "Mb")
        print(type(Memory), type(mega))
        self.assertEqual(mega.to("byte"), 1024**2)
        self.assertEqual(mega, Memory(1, "mb"))

        same_mega = Memory.from_string("1Mb")
        self.assertEqual(same_mega.unit_type, "memory")

        other_mega = Memory.from_string("+1.0 mb")
        self.assertEqual(mega, other_mega)

    def test_unitized(self):

        @unitized("eV")
        def f():
            return [1, 2, 3]

        self.assertEqual(str(f()[0]), "1.0 eV")
        self.assertIsInstance(f(), list)

        @unitized("eV")
        def g():
            return 2, 3, 4

        self.assertEqual(str(g()[0]), "2.0 eV")
        self.assertIsInstance(g(), tuple)

        @unitized("pm")
        def h():
            d = collections.OrderedDict()
            for i in range(3):
                d[i] = i * 20
            return d

        self.assertEqual(str(h()[1]), "20.0 pm")
        self.assertIsInstance(h(), collections.OrderedDict)

        @unitized("kg")
        def i():
            return FloatWithUnit(5, "g")

        self.assertEqual(i(), FloatWithUnit(0.005, "kg"))

        @unitized("kg")
        def j():
            return ArrayWithUnit([5, 10], "g")

        j_out = j()
        self.assertEqual(j_out.unit, Unit("kg"))
        self.assertEqual(j_out[0], 0.005)
        self.assertEqual(j_out[1], 0.01)

    def test_compound_operations(self):
        g = 10 * Length(1, "m") / (Time(1, "s") ** 2)
        e = Mass(1, "kg") * g * Length(1, "m")
        self.assertEqual(str(e), "10.0 N m")
        form_e = FloatWithUnit(10, unit="kJ mol^-1").to("eV atom^-1")
        self.assertAlmostEqual(float(form_e), 0.103642691905)
        self.assertEqual(str(form_e.unit), "eV atom^-1")
        self.assertRaises(UnitError, form_e.to, "m s^-1")
        a = FloatWithUnit(1.0, "Ha^3")
        b = a.to("J^3")
        self.assertAlmostEqual(b, 8.28672661615e-53)
        self.assertEqual(str(b.unit), "J^3")
        a = FloatWithUnit(1.0, "Ha bohr^-2")
        b = a.to("J m^-2")
        self.assertAlmostEqual(b, 1556.8929672799854)
        self.assertEqual(str(b.unit), "J m^-2")

    def test_as_base_units(self):
        x = FloatWithUnit(5, "MPa")
        self.assert_equal(FloatWithUnit(5000000, "Pa"), x.as_base_units)

class ArrayWithFloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        """
        Similar to FloatWithUnitTest.test_energy.
        Check whether EnergyArray and FloatWithUnit have same behavior.

        # TODO
        One can merge the two tests easily:

        for obj in [Energy, EnergyArray]:
            a = obj(...)
            self.assert(...)

        """
        a = EnergyArray(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(float(b), 0.0404242579378)
        c = EnergyArray(3.14, "J")
        self.assertAlmostEqual(float(c.to("eV")), 1.9598339337836966e+19)
        # self.assertRaises(ValueError, Energy, 1, "m")

        d = EnergyArray(1, "Ha")
        self.assertAlmostEqual(float(a + d), 28.311385050630882)
        self.assertAlmostEqual(float(a - d), -26.11138505063088)
        self.assertEqual(float(a + 1), 2.1)

    def test_time(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        # here there's a minor difference because we have a ndarray with
        # dtype=np.int.
        a = TimeArray(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60 h")
        self.assertEqual(str(3 * a), "60 h")

    def test_length(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        x = LengthArray(4.2, "ang")
        self.assertAlmostEqual(float(x.to("cm")), 4.2e-08)
        self.assertEqual(float(x.to("pm")), 420)
        self.assertEqual(str(x / 2), "2.1 ang")

    def test_array_algebra(self):
        ene_ha = EnergyArray([1, 2], "Ha")
        ene_ev = EnergyArray([1, 2], "eV")
        time_s = TimeArray([1, 2], "s")

        e1 = ene_ha.copy()
        e1 += 1
        e2 = ene_ha.copy()
        e2 -= 1
        e3 = ene_ha.copy()
        #e3 /= 2
        e4 = ene_ha.copy()
        e4 *= 2

        objects_with_unit = [
            ene_ha + ene_ev,
            ene_ha - ene_ev,
            3 * ene_ha,
            ene_ha * 3,
            ene_ha / 3,
            3 / ene_ha,
            ene_ha * time_s,
            ene_ha / ene_ev,
            ene_ha.copy(),
            ene_ha[0:1],
            e1,
            e2,
            e3,
            e4,
        ]

        for i, obj in enumerate(objects_with_unit):
            #print(i, obj.unit)
            self.assertTrue(hasattr(obj, "unit"))
            #self.assertTrue(str(obj.unit) == "Ha")

        objects_without_unit = [
            # Here we could return a FloatWithUnit object but I prefer this
            # a bare scalar since FloatWithUnit extends float while we could
            # have an int.
            ene_ha[0],
        ]

        for obj in objects_without_unit:
            self.assertFalse(hasattr(obj, "unit"))

        with self.assertRaises(UnitError):
            ene_ha + time_s

    def test_factors(self):
        e = EnergyArray([27.21138386, 1], "eV").to("Ha")
        self.assertTrue(str(e) == "[ 0.99999996  0.03674932] Ha")
        l = LengthArray([1.0], "ang").to("bohr")
        self.assertTrue(str(l) == "[ 1.88972612] bohr")
        v = ArrayWithUnit([1, 2, 3], "bohr^3").to("ang^3")
        self.assertTrue(str(v) == '[ 0.14818471  0.29636942  0.44455413] ang^3')

    def test_as_base_units(self):
        x = ArrayWithUnit([5, 10], "MPa")
        self.assert_equal(ArrayWithUnit([5000000, 10000000], "Pa"), x.as_base_units)

class DataPersistenceTest(PymatgenTest):
    def test_pickle(self):
        """Test whether FloatWithUnit and ArrayWithUnit support pickle"""

        for cls in [FloatWithUnit, ArrayWithUnit]:
            a = cls(1, "eV")
            b = cls(10, "N bohr")
            objects = [a, b]

            new_objects_from_protocol = self.serialize_with_pickle(objects)

            for new_objects in new_objects_from_protocol:
                for old_item, new_item in zip(objects, new_objects):
                    self.assertTrue(str(old_item) == str(new_item))

if __name__ == '__main__':
    import unittest
    unittest.main()
<|endoftext|>


<commit_before>import subprocess, os, signal, time, sys

def cmd_join(x):
    """The opposite of shlex.split()."""
    y = []
    for e in x:
        if " " not in e:
            y.append(e)
        else:
            y.append("\"%s\"" % e)
    return " ".join(y)

def signal_with_number(number):
    """Converts a numeric signal ID to a string like "SIGINT"."""
    for name in dir(signal):
        if name.startswith("SIG") and getattr(signal, name) == number:
            return name
    return str(number)

def wait_with_timeout(process, timeout):
    if timeout is None:
        return process.wait()
    else:
        wait_interval = 0.05
        n_intervals = int(timeout / wait_interval) + 1
        while process.poll() is None:
            n_intervals -= 1
            if n_intervals == 0: break
            time.sleep(wait_interval)
        return process.poll()

valgrind_error_code = 100

class SubProcess(object):
    def __init__(self, command_line, output_path, test_dir, valgrind_tool=None, interactive=False):
        if not os.path.exists(command_line[0]):
            raise RuntimeError("Bad executable: %r does not exist." % command_line[0])

        self.summary = os.path.basename(command_line[0])
        if len(command_line) > 4: self.summary += " " + cmd_join(command_line[1:4]) + " [...]"
        else: self.summary += " " + cmd_join(command_line[1:])

        self.valgrind_tool = valgrind_tool
        if valgrind_tool is not None:
            cmd_line = ["valgrind",
                "--log-file=%s" % test_dir.make_file(output_path + ".valgrind"),
                "--tool=%s" % valgrind_tool,
                "--error-exitcode=%d" % valgrind_error_code]
            if valgrind_tool == "memcheck":
                cmd_line.append("--leak-check=full")
                cmd_line.append("--track-origins=yes")
                cmd_line.append("--show-reachable=yes")
            if valgrind_tool == "drd":
                cmd_line.append("--read-var-info=yes")
            if interactive: cmd_line.append("--db-attach=yes")
            command_line = cmd_line + command_line

        self.interactive = interactive
        if interactive:
            self.output_path = "<stdout>"
            subprocess_stdout = sys.stdout
        else:
            self.output_path = test_dir.make_file(output_path)
            subprocess_stdout = file(self.output_path, "w")

        self.sub = subprocess.Popen(command_line,
            stdout=subprocess_stdout, stderr=subprocess.STDOUT)

    def verify(self):
        res = self.sub.poll()
        if res is None:
            return   # All is well

        if res == 0:
            raise RuntimeError("%r shut down without being asked to, but return code was 0." % \
                self.summary)
        elif self.valgrind_tool is not None and res == valgrind_error_code:
            raise RuntimeError("%r terminated unexpectedly with Valgrind errors. Output is at %r." % \
                (self.summary, self.output_path))
        elif res > 0:
            raise RuntimeError("%r terminated unexpectedly with return code %d. Output is at %r." % \
                (self.summary, res, self.output_path))
        elif res < 0:
            raise RuntimeError("%r terminated unexpectedly with signal %s. Output is at %r." % \
                (self.summary, signal_with_number(-res), self.output_path))

    def interrupt(self, timeout = 5):
        self.verify()

        try: self.sub.send_signal(signal.SIGINT)
        except: pass

        if self.valgrind_tool is not None:
            # Sometimes valgrind loses track of the first SIGINT, especially if it receives it
            # right as it is starting up. Send a second SIGINT to make sure the server gets it.
            # This is stupid and we shouldn't have to do this.
            if wait_with_timeout(self.sub, 10) is None:
                self.sub.send_signal(signal.SIGINT)

        self.wait(timeout)

    def wait(self, timeout):
        try:
            res = wait_with_timeout(self.sub, timeout)
        finally:
            try: sub.terminate()
            except: pass

        if res is None:
            raise RuntimeError("%r took more than %d seconds to shut down. Output is at %r." % \
                (self.summary, timeout, self.output_path))
        elif self.valgrind_tool is not None and res == valgrind_error_code:
            raise RuntimeError("Valgrind reported errors in %r. Output is at %r." % \
                (self.summary, self.output_path))
        elif res > 0:
            raise RuntimeError("%r exited with exit status %d. Output is at %r." % \
                (self.summary, res, self.output_path))
        elif res < 0:
            raise RuntimeError("%r exited with signal %s. Output is at %r." % \
                (self.summary, signal_with_number(-res), self.output_path))

    def kill(self):
        self.verify()
        try: self.sub.terminate()
        except: pass

    def __del__(self):
        # Avoid leaking processes in case there is an error
        if hasattr(self, "sub") and self.sub.poll() is None:
            self.sub.send_signal(signal.SIGKILL)

def run_executable(command_line, output_path, test_dir, timeout=60, valgrind_tool=None):

    summary = os.path.basename(command_line[0])
    for chunk in command_line[1:]:
        if len(summary) + len(chunk) > 200:
            summary += " [...]"
            break
        elif " " in chunk:
            summary += " \"%s\"" % chunk
        else:
            summary += " " + chunk
    print "Running %r;" % summary,
    sys.stdout.flush()

    sub = SubProcess(command_line, output_path, test_dir, valgrind_tool)

    start_time = time.time()
    sub.wait(timeout)
    end_time = time.time()

    print "done after %f seconds." % (end_time - start_time)
<commit_msg>Fixed undefined variable problem in procutil.py<commit_after>import subprocess, os, signal, time, sys

def cmd_join(x):
    """The opposite of shlex.split()."""
    y = []
    for e in x:
        if " " not in e:
            y.append(e)
        else:
            y.append("\"%s\"" % e)
    return " ".join(y)

def signal_with_number(number):
    """Converts a numeric signal ID to a string like "SIGINT"."""
    for name in dir(signal):
        if name.startswith("SIG") and getattr(signal, name) == number:
            return name
    return str(number)

def wait_with_timeout(process, timeout):
    if timeout is None:
        return process.wait()
    else:
        wait_interval = 0.05
        n_intervals = int(timeout / wait_interval) + 1
        while process.poll() is None:
            n_intervals -= 1
            if n_intervals == 0: break
            time.sleep(wait_interval)
        return process.poll()

valgrind_error_code = 100

class SubProcess(object):
    def __init__(self, command_line, output_path, test_dir, valgrind_tool=None, interactive=False):
        if not os.path.exists(command_line[0]):
            raise RuntimeError("Bad executable: %r does not exist." % command_line[0])

        self.summary = os.path.basename(command_line[0])
        if len(command_line) > 4: self.summary += " " + cmd_join(command_line[1:4]) + " [...]"
        else: self.summary += " " + cmd_join(command_line[1:])

        self.valgrind_tool = valgrind_tool
        if valgrind_tool is not None:
            cmd_line = ["valgrind",
                "--log-file=%s" % test_dir.make_file(output_path + ".valgrind"),
                "--tool=%s" % valgrind_tool,
                "--error-exitcode=%d" % valgrind_error_code]
            if valgrind_tool == "memcheck":
                cmd_line.append("--leak-check=full")
                cmd_line.append("--track-origins=yes")
                cmd_line.append("--show-reachable=yes")
            if valgrind_tool == "drd":
                cmd_line.append("--read-var-info=yes")
            if interactive: cmd_line.append("--db-attach=yes")
            command_line = cmd_line + command_line

        self.interactive = interactive
        if interactive:
            self.output_path = "<stdout>"
            subprocess_stdout = sys.stdout
        else:
            self.output_path = test_dir.make_file(output_path)
            subprocess_stdout = file(self.output_path, "w")

        self.sub = subprocess.Popen(command_line,
            stdout=subprocess_stdout, stderr=subprocess.STDOUT)

    def verify(self):
        res = self.sub.poll()
        if res is None:
            return   # All is well

        if res == 0:
            raise RuntimeError("%r shut down without being asked to, but return code was 0." % \
                self.summary)
        elif self.valgrind_tool is not None and res == valgrind_error_code:
            raise RuntimeError("%r terminated unexpectedly with Valgrind errors. Output is at %r." % \
                (self.summary, self.output_path))
        elif res > 0:
            raise RuntimeError("%r terminated unexpectedly with return code %d. Output is at %r." % \
                (self.summary, res, self.output_path))
        elif res < 0:
            raise RuntimeError("%r terminated unexpectedly with signal %s. Output is at %r." % \
                (self.summary, signal_with_number(-res), self.output_path))

    def interrupt(self, timeout = 5):
        self.verify()

        try: self.sub.send_signal(signal.SIGINT)
        except: pass

        if self.valgrind_tool is not None:
            # Sometimes valgrind loses track of the first SIGINT, especially if it receives it
            # right as it is starting up. Send a second SIGINT to make sure the server gets it.
            # This is stupid and we shouldn't have to do this.
            if wait_with_timeout(self.sub, 10) is None:
                self.sub.send_signal(signal.SIGINT)

        self.wait(timeout)

    def wait(self, timeout):
        try:
            res = wait_with_timeout(self.sub, timeout)
        finally:
            try:
                self.sub.terminate()
            except:
                pass

        if res is None:
            raise RuntimeError("%r took more than %d seconds to shut down. Output is at %r." % \
                (self.summary, timeout, self.output_path))
        elif self.valgrind_tool is not None and res == valgrind_error_code:
            raise RuntimeError("Valgrind reported errors in %r. Output is at %r." % \
                (self.summary, self.output_path))
        elif res > 0:
            raise RuntimeError("%r exited with exit status %d. Output is at %r." % \
                (self.summary, res, self.output_path))
        elif res < 0:
            raise RuntimeError("%r exited with signal %s. Output is at %r." % \
                (self.summary, signal_with_number(-res), self.output_path))

    def kill(self):
        self.verify()
        try: self.sub.terminate()
        except: pass

    def __del__(self):
        # Avoid leaking processes in case there is an error
        if hasattr(self, "sub") and self.sub.poll() is None:
            self.sub.send_signal(signal.SIGKILL)

def run_executable(command_line, output_path, test_dir, timeout=60, valgrind_tool=None):

    summary = os.path.basename(command_line[0])
    for chunk in command_line[1:]:
        if len(summary) + len(chunk) > 200:
            summary += " [...]"
            break
        elif " " in chunk:
            summary += " \"%s\"" % chunk
        else:
            summary += " " + chunk
    print "Running %r;" % summary,
    sys.stdout.flush()

    sub = SubProcess(command_line, output_path, test_dir, valgrind_tool)

    start_time = time.time()
    sub.wait(timeout)
    end_time = time.time()

    print "done after %f seconds." % (end_time - start_time)
<|endoftext|>


<commit_before>import json
from datetime import datetime, date, timedelta
from wq.io import JsonNetIO, BaseIO, TimeSeriesMapper, make_date_mapper
from .constants import *

parse_date = make_date_mapper('%Y-%m-%d')

class AcisIO(JsonNetIO):
    """
    Base class for loading data from ACIS web services
    See http://data.rcc-acis.org/doc/
    """

    path = None  # ACIS web service path
    basin = None  # HUC8 watershed
    elems = None  # Requested element codes
    meta = None  # Requested metadata

    @property
    def url(self):
        """
        URL for wq.io.loaders.NetLoader
        """
        return "http://data.rcc-acis.org/%s" % self.path

    def getdate(self, name):
        """
        Retrieve given property from class/instance, ensuring it is a date
        """
        value = getattr(self, name, None)

        if value is None:
            return None

        if isinstance(value, basestring):
            value = parse_date(value)

        if isinstance(value, datetime):
            value = value.date()

        return value

    def getlist(self, name):
        """
        Retrieve given property from class/instance, ensuring it is a list.
        Also determine whether the list contains simple text/numeric values or
        nested dictionaries (a "complex" list)
        """
        value = getattr(self, name, None)
        complex = False

        def str_value(val):
            if isinstance(val, dict):
                complex = True
                return val
            else:
                return unicode(val)

        if value is None:
            pass
        elif isinstance(value, list):
            value = [str_value(val) for val in value]
        else:
            value = [str_value(value)]

        return value, complex

    def set_param(self, into, name):
        """
        Set parameter key, noting whether list value is "complex"
        """
        value, complex = self.getlist(name)
        if value is not None:
            into[name] = value
        return complex

    def get_params(self):
        """
        Get parameters for web service, noting whether any are "complex"
        """
        params = {}
        complex = False

        if self.basin is None:
            # FIXME: Support other region types
            raise NotImplementedError("No region specified")

        for param in ('basin', 'meta', 'elems', 'sdate', 'edate'):
            if self.set_param(params, param):
                complex = True
        return params, complex

    @property
    def params(self):
        """
        URL parameters for wq.io.loaders.NetLoader
        """
        params, complex = self.get_params()
        if complex:
            # ACIS web service supports JSON object as "params" parameter
            params = {
                key: val[0]
                if len(val) == 1 and isinstance(val[0], basestring)
                else val
                for key, val in params.items()
            }
            return {'params': json.dumps(params)}
        else:
            # Simpler queries can use traditional URL parameters
            return {
                key: ','.join(val)
                for key, val in params.items()
            }

class StationMetaIO(AcisIO):
    """
    Retrieves metadata about the climate stations in a region.
    See http://data.rcc-acis.org/doc/#title8
    """

    namespace = "meta"  # For wq.io.parsers.text.JsonParser
    path = "StnMeta"
    elems = None

    @property
    def meta(self):
        """
        Request all available metadata by default
        """
        meta = ['name', 'state', 'sids', 'll',
                'elev', 'uid', 'county', 'climdiv']
        if self.elems is not None:
            meta.append('valid_daterange')
        return meta

    def map_value(self, field, value):
        """
        Clean up some values returned from the web service.
        (overrides wq.io.mappers.BaseMapper)
        """

        if field == 'sids':
            # Site identifiers are returned as "[id] [auth_id]";
            # Map to auth name for easier usability
            ids = {}
            for idinfo in value:
                id, auth = idinfo.split(' ')
                auth = AUTHORITY_BY_ID[auth]
                ids[auth['name']] = id
            return ids

        if field == 'valid_daterange':
            # Date ranges for each element are returned in an array
            # (sorted by the order the elements were were requested);
            # Convert to dictionary with element id as key
            elems, complex = self.getlist('elems')
            ranges = {}
            for elem, val in zip(elems, value):
                if val:
                    start, end = val
                    ranges[elem] = (parse_date(start), parse_date(end))
                else:
                    ranges[elem] = None, None
            return ranges
        return value

class StationDataIO(StationMetaIO):
    """
    Retrieve daily time series data from the climate stations in a region.
    See http://data.rcc-acis.org/doc/#title19
    """

    namespace = "data"  # For wq.io.parsers.text.JsonParser
    path = "MultiStnData"

    sdate = None  # Start date
    edate = None  # End date
    add = None  # Additional information for daily results

    def get_field_names(self):
        """
        ACIS web service returns "meta" and "data" for each station;
        Use meta attributes as field names
        """
        field_names = super(StationDataIO, self).get_field_names()
        if field_names == ['meta', 'data']:
            field_names = self.data[0]['meta'].keys()
        return field_names

    def get_params(self):
        if self.sdate is None or self.edate is None:
            raise NotImplementedError("sdate and edate must be set!")
        params, complex = super(StationDataIO, self).get_params()

        # If self.add is defined, set it for each requested element
        if self.add is not None:
            add, add_is_complex = self.getlist('add')
            complex = True
            elems = []
            for elem in params.get('elems', []):
                if not isinstance(elem, dict):
                    elem = {'name': elem}
                elem['add'] = ",".join(add)
                elems.append(elem)
            params['elems'] = elems

        return params, complex

    def usable_item(self, data):
        """
        ACIS web service returns "meta" and "data" for each station; use meta
        attributes as item values, and add an IO for iterating over "data"
        """

        # Use metadata as item
        item = data['meta']

        # Add nested IO for data
        elems, elems_is_complex = self.getlist('elems')
        if elems_is_complex:
            elems = [elem['name'] for elem in elems]

        add, add_is_complex = self.getlist('add')
        item['data'] = DataIO(
            data=data['data'],
            elems=elems,
            add=add,
            sdate=self.getdate('sdate'),
            edate=self.getdate('edate')
        )

        # TupleMapper will convert item to namedtuple
        return super(StationDataIO, self).usable_item(item)

class DataIO(TimeSeriesMapper, BaseIO):
    """
    IO for iterating over ACIS time series data.
    Created internally by StationDataIO; not meant to be used directly.
    """

    # Inherited from parent
    elems = []
    add = []
    sdate = None
    edate = None

    date_formats = []  # For TimeSeriesMapper

    def load_data(self, data):
        """
        MultiStnData data results are arrays without explicit dates;
        Infer time series based on start date.
        """

        date = self.sdate
        day = timedelta(days=1)
        for row in data:
            data = {'date': date}
            if self.add:
                # If self.add is set, results will contain additional
                # attributes (e.g. flags). In that case, create one row per
                # result, with attributes "date", "elem", "value", and one for
                # each item in self.add.
                for elem, vals in zip(self.elems, row):
                    data['elem'] = elem
                    for add, val in zip(['value'] + self.add, vals):
                        data[add] = val
                    yield data
            else:
                # Otherwise, return one row per date, with "date" and each
                # element's value as attributes.
                for elem, val in zip(self.elems, row):
                    # namedtuple doesn't like numeric field names
                    if elem.isnumeric():
                        elem = "e%s" % elem
                    data[elem] = val
                yield data
            date += day

    def __init__(self, *args, **kwargs):
        data = kwargs.pop('data')
        self.data = self.load_data(data)
        super(DataIO, self).__init__(*args, **kwargs)

    def get_field_names(self):
        """
        Different field names depending on self.add setting (see load_data)
        For BaseIO
        """
        if self.add:
            return ['date', 'elem', 'value'] + [flag for flag in self.add]
        else:
            field_names = ['date']
            for elem in self.elems:
                # namedtuple doesn't like numeric field names
                if elem.isnumeric():
                    elem = "e%s" % elem
                field_names.append(elem)
            return field_names

    @property
    def key_fields(self):
        """
        Different key fields depending on self.add setting (see load_data)
        For TimeSeriesMapper
        """
        if self.add:
            return ['date', 'elem']
        else:
            return ['date']
<commit_msg>handle dates differently<commit_after>import json
from datetime import datetime, date, timedelta
from wq.io import JsonNetIO, BaseIO, TimeSeriesMapper, make_date_mapper
from .constants import *

parse_date = make_date_mapper('%Y-%m-%d')

class AcisIO(JsonNetIO):
    """
    Base class for loading data from ACIS web services
    See http://data.rcc-acis.org/doc/
    """

    path = None  # ACIS web service path
    basin = None  # HUC8 watershed
    elems = None  # Requested element codes
    meta = None  # Requested metadata

    @property
    def url(self):
        """
        URL for wq.io.loaders.NetLoader
        """
        return "http://data.rcc-acis.org/%s" % self.path

    def getdate(self, name):
        """
        Retrieve given property from class/instance, ensuring it is a date
        """
        value = getattr(self, name, None)

        if value is None:
            return None

        if isinstance(value, basestring):
            value = parse_date(value)

        if isinstance(value, datetime):
            value = value.date()

        return value

    def getlist(self, name):
        """
        Retrieve given property from class/instance, ensuring it is a list.
        Also determine whether the list contains simple text/numeric values or
        nested dictionaries (a "complex" list)
        """
        value = getattr(self, name, None)
        complex = False

        def str_value(val):
            if isinstance(val, dict):
                complex = True
                return val
            else:
                return unicode(val)

        if value is None:
            pass
        elif isinstance(value, list):
            value = [str_value(val) for val in value]
        else:
            value = [str_value(value)]

        return value, complex

    def set_param(self, into, name):
        """
        Set parameter key, noting whether list value is "complex"
        """
        value, complex = self.getlist(name)
        if value is not None:
            into[name] = value
        return complex

    def get_params(self):
        """
        Get parameters for web service, noting whether any are "complex"
        """
        params = {}
        complex = False

        if self.basin is None:
            # FIXME: Support other region types
            raise NotImplementedError("No region specified")

        for param in ('basin', 'meta', 'elems'):
            if self.set_param(params, param):
                complex = True
        for param in ('sdate', 'edate'):
            value = self.getdate(param)
            if value:
                params[param] = [unicode(value)]
        return params, complex

    @property
    def params(self):
        """
        URL parameters for wq.io.loaders.NetLoader
        """
        params, complex = self.get_params()
        if complex:
            # ACIS web service supports JSON object as "params" parameter
            params = {
                key: val[0]
                if len(val) == 1 and isinstance(val[0], basestring)
                else val
                for key, val in params.items()
            }
            return {'params': json.dumps(params)}
        else:
            # Simpler queries can use traditional URL parameters
            return {
                key: ','.join(val)
                for key, val in params.items()
            }

class StationMetaIO(AcisIO):
    """
    Retrieves metadata about the climate stations in a region.
    See http://data.rcc-acis.org/doc/#title8
    """

    namespace = "meta"  # For wq.io.parsers.text.JsonParser
    path = "StnMeta"
    elems = None

    @property
    def meta(self):
        """
        Request all available metadata by default
        """
        meta = ['name', 'state', 'sids', 'll',
                'elev', 'uid', 'county', 'climdiv']
        if self.elems is not None:
            meta.append('valid_daterange')
        return meta

    def map_value(self, field, value):
        """
        Clean up some values returned from the web service.
        (overrides wq.io.mappers.BaseMapper)
        """

        if field == 'sids':
            # Site identifiers are returned as "[id] [auth_id]";
            # Map to auth name for easier usability
            ids = {}
            for idinfo in value:
                id, auth = idinfo.split(' ')
                auth = AUTHORITY_BY_ID[auth]
                ids[auth['name']] = id
            return ids

        if field == 'valid_daterange':
            # Date ranges for each element are returned in an array
            # (sorted by the order the elements were were requested);
            # Convert to dictionary with element id as key
            elems, complex = self.getlist('elems')
            ranges = {}
            for elem, val in zip(elems, value):
                if val:
                    start, end = val
                    ranges[elem] = (parse_date(start), parse_date(end))
                else:
                    ranges[elem] = None, None
            return ranges
        return value

class StationDataIO(StationMetaIO):
    """
    Retrieve daily time series data from the climate stations in a region.
    See http://data.rcc-acis.org/doc/#title19
    """

    namespace = "data"  # For wq.io.parsers.text.JsonParser
    path = "MultiStnData"

    sdate = None  # Start date
    edate = None  # End date
    add = None  # Additional information for daily results

    def get_field_names(self):
        """
        ACIS web service returns "meta" and "data" for each station;
        Use meta attributes as field names
        """
        field_names = super(StationDataIO, self).get_field_names()
        if field_names == ['meta', 'data']:
            field_names = self.data[0]['meta'].keys()
        return field_names

    def get_params(self):
        if self.sdate is None or self.edate is None:
            raise NotImplementedError("sdate and edate must be set!")
        params, complex = super(StationDataIO, self).get_params()

        # If self.add is defined, set it for each requested element
        if self.add is not None:
            add, add_is_complex = self.getlist('add')
            complex = True
            elems = []
            for elem in params.get('elems', []):
                if not isinstance(elem, dict):
                    elem = {'name': elem}
                elem['add'] = ",".join(add)
                elems.append(elem)
            params['elems'] = elems

        return params, complex

    def usable_item(self, data):
        """
        ACIS web service returns "meta" and "data" for each station; use meta
        attributes as item values, and add an IO for iterating over "data"
        """

        # Use metadata as item
        item = data['meta']

        # Add nested IO for data
        elems, elems_is_complex = self.getlist('elems')
        if elems_is_complex:
            elems = [elem['name'] for elem in elems]

        add, add_is_complex = self.getlist('add')
        item['data'] = DataIO(
            data=data['data'],
            elems=elems,
            add=add,
            sdate=self.getdate('sdate'),
            edate=self.getdate('edate')
        )

        # TupleMapper will convert item to namedtuple
        return super(StationDataIO, self).usable_item(item)

class DataIO(TimeSeriesMapper, BaseIO):
    """
    IO for iterating over ACIS time series data.
    Created internally by StationDataIO; not meant to be used directly.
    """

    # Inherited from parent
    elems = []
    add = []
    sdate = None
    edate = None

    date_formats = []  # For TimeSeriesMapper

    def load_data(self, data):
        """
        MultiStnData data results are arrays without explicit dates;
        Infer time series based on start date.
        """

        date = self.sdate
        day = timedelta(days=1)
        for row in data:
            data = {'date': date}
            if self.add:
                # If self.add is set, results will contain additional
                # attributes (e.g. flags). In that case, create one row per
                # result, with attributes "date", "elem", "value", and one for
                # each item in self.add.
                for elem, vals in zip(self.elems, row):
                    data['elem'] = elem
                    for add, val in zip(['value'] + self.add, vals):
                        data[add] = val
                    yield data
            else:
                # Otherwise, return one row per date, with "date" and each
                # element's value as attributes.
                for elem, val in zip(self.elems, row):
                    # namedtuple doesn't like numeric field names
                    if elem.isnumeric():
                        elem = "e%s" % elem
                    data[elem] = val
                yield data
            date += day

    def __init__(self, *args, **kwargs):
        data = kwargs.pop('data')
        self.data = self.load_data(data)
        super(DataIO, self).__init__(*args, **kwargs)

    def get_field_names(self):
        """
        Different field names depending on self.add setting (see load_data)
        For BaseIO
        """
        if self.add:
            return ['date', 'elem', 'value'] + [flag for flag in self.add]
        else:
            field_names = ['date']
            for elem in self.elems:
                # namedtuple doesn't like numeric field names
                if elem.isnumeric():
                    elem = "e%s" % elem
                field_names.append(elem)
            return field_names

    @property
    def key_fields(self):
        """
        Different key fields depending on self.add setting (see load_data)
        For TimeSeriesMapper
        """
        if self.add:
            return ['date', 'elem']
        else:
            return ['date']
<|endoftext|>


<commit_before># Author: Immanuel Bayer
# License: BSD 3 clause

import numpy as np
from ffm import ffm_predict
import scipy.sparse as sp
from sklearn.utils import check_random_state
from sklearn.metrics import mean_squared_error, r2_score

def make_user_item_regression(random_state=123, n_user=20, n_item=20,
                              label_stdev=0.4, rank=2, bias=True,
                              first_order=True, stdev_w0=.2, stdev_w=0.3,
                              stdev_V=0.4, mean_w0=2, mean_w=5, mean_V=10):

    n_features = n_user + n_item
    n_samples = n_user * n_item
    # create design matrix
    user_cols = np.repeat(range(n_user), n_item)
    item_cols = np.array(range(n_item) * n_user) + n_user
    cols = np.hstack((user_cols, item_cols))
    rows = np.hstack((np.arange(n_item*n_user), np.arange(n_item*n_user)))

    X = sp.coo_matrix((np.ones_like(cols, dtype=np.float64), (rows, cols)))
    X = sp.csc_matrix(X)
    assert X.shape[0] == n_samples
    assert X.shape[1] == n_features

    # sample the model parameter
    random_state = check_random_state(random_state)
    w0 = random_state.normal(mean_w0, stdev_w0)
    w = random_state.normal(mean_w, stdev_w, n_features)
    V = random_state.normal(mean_V, stdev_V, (rank, n_features))

    y = ffm_predict(w0, w, V, X)
    if label_stdev > 0:
        y = random_state.normal(y, label_stdev)

    return X, y, (w0, w, V)

if __name__ == '__main__':
    X, y, coef = make_user_item_regression(n_user=5, n_item=5, rank=2,
                                           label_stdev=2)
    from sklearn.cross_validation import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.33, random_state=42)

    from mcmc import FMRegression
    fm = FMRegression(rank=2)
    y_pred = fm.fit_predict(sp.csc_matrix(X_train), y_train,
                            sp.csc_matrix(X_test))

    print('rmse', mean_squared_error(y_pred, y_test))
    print('r2_score', r2_score(y_pred, y_test))
    np.random.shuffle(y_pred)
    print('----  shuffled pred ---------')
    print('rmse', mean_squared_error(y_pred, y_test))
    print('r2_score', r2_score(y_pred, y_test))
<commit_msg>Wrap range() by list() for python3<commit_after># Author: Immanuel Bayer
# License: BSD 3 clause

import numpy as np
from ffm import ffm_predict
import scipy.sparse as sp
from sklearn.utils import check_random_state
from sklearn.metrics import mean_squared_error, r2_score

def make_user_item_regression(random_state=123, n_user=20, n_item=20,
                              label_stdev=0.4, rank=2, bias=True,
                              first_order=True, stdev_w0=.2, stdev_w=0.3,
                              stdev_V=0.4, mean_w0=2, mean_w=5, mean_V=10):

    n_features = n_user + n_item
    n_samples = n_user * n_item
    # create design matrix
    user_cols = np.repeat(range(n_user), n_item)
    item_cols = np.array(list(range(n_item)) * n_user) + n_user
    cols = np.hstack((user_cols, item_cols))
    rows = np.hstack((np.arange(n_item*n_user), np.arange(n_item*n_user)))

    X = sp.coo_matrix((np.ones_like(cols, dtype=np.float64), (rows, cols)))
    X = sp.csc_matrix(X)
    assert X.shape[0] == n_samples
    assert X.shape[1] == n_features

    # sample the model parameter
    random_state = check_random_state(random_state)
    w0 = random_state.normal(mean_w0, stdev_w0)
    w = random_state.normal(mean_w, stdev_w, n_features)
    V = random_state.normal(mean_V, stdev_V, (rank, n_features))

    y = ffm_predict(w0, w, V, X)
    if label_stdev > 0:
        y = random_state.normal(y, label_stdev)

    return X, y, (w0, w, V)

if __name__ == '__main__':
    X, y, coef = make_user_item_regression(n_user=5, n_item=5, rank=2,
                                           label_stdev=2)
    from sklearn.cross_validation import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.33, random_state=42)

    from mcmc import FMRegression
    fm = FMRegression(rank=2)
    y_pred = fm.fit_predict(sp.csc_matrix(X_train), y_train,
                            sp.csc_matrix(X_test))

    print('rmse', mean_squared_error(y_pred, y_test))
    print('r2_score', r2_score(y_pred, y_test))
    np.random.shuffle(y_pred)
    print('----  shuffled pred ---------')
    print('rmse', mean_squared_error(y_pred, y_test))
    print('r2_score', r2_score(y_pred, y_test))
<|endoftext|>


<commit_before>import os

from django.http import HttpResponse, HttpResponseForbidden, JsonResponse
from django.views.decorators.csrf import csrf_exempt
from django.views.generic import View

from .models import Memory

SLACK_TOKEN = os.environ['SLACK_TOKEN']
BOT_NAME = 'membot'
KNOWN_COMMANDS = ['show',]

def homepage(request):
    return HttpResponse('Hello world this is membot')

class CommandView(View):
    response = {'text': ''}

    @csrf_exempt
    def dispatch(self, *args, **kwargs):
        return super(CommandView, self).dispatch(*args, **kwargs)

    def post(self, request, *args, **kwargs):
        received = request.POST
        auth_token = received.get('token', None)

        # make sure command is coming from the right place
        if auth_token != SLACK_TOKEN:
            return HttpResponseForbidden()

        # collect our command params
        command = self.parse_command(received)

        # if we don't actually have a command ...
        if not command['text']:
            self.set_response('Yes <@%s>?' % (command['person']))

        # otherwise take action
        else:
            if 'special' in command:
                action = command['special']

                if action not in KNOWN_COMMANDS:
                    self.set_response('\'%s\' sounds like a special command, <@%s>, but I haven\'t learned that one yet :(' % (command['special'], command['person']))

                if action == 'show':
                    memories = []
                    #TODO: handle `since` option

                    memories = Memory.objects.filter(is_active=True, category__in=command['categories']).distinct()

                    if memories:
                        intro = ['Here\'s what I remember about %s:\n' % ' '.join(command['categories'])]
                        for memory in memories:
                            say = '- On %s, <@%s> said: %s' % (memory.created.strftime('%B %d'), memory.person, memory.text)
                            report.append(say)
                        self.set_response(intro + '```' + '\n'.join(report) + '```')
                    else:
                        self.set_response('Sorry, I don\'t remember anything like that!')

            # we have a memory to log, so do it for each defined category
            else:
                for category in command['categories']:
                    kwargs = {
                        'text': command['text'],
                        'person': command['person'],
                        'category': category,
                    }
                    memory = Memory(**kwargs)
                    memory.save()

                self.set_response('Got it, <@%s>!' % (command['person']))

        return JsonResponse(self.response)

    def set_response(self, text):
        self.response['text'] = text

    def parse_command(self, received):
        # split the command text for cleaning
        tokens = received.get('text', '').split(' ')

        # we don't need the bot name, and using `startswith`
        # catches most natural language punctuation
        if tokens[0].lower().startswith(BOT_NAME):
            del tokens[0]

        # just in case we had punctuation after the trigger name
        if tokens and tokens[0].lower() in [',', ':', '!', '?']:
            del tokens[0]

        # collect the hashtags
        categories = []
        for token in tokens:
            if token.startswith('#'):
                categories.append(token.lower())

        # give ourselves a default category if nothing else
        if not categories:
            categories = ['#general']

        # begin the command dict
        command = {
            'person': received.get('user_name', None),
            'categories': categories,
        }

        # see if we have a special command
        if tokens and tokens[0].lower() in ['please', 'plz']:
            del tokens[0]

            # take the next word for our special command
            if tokens:
                command.update({
                    'special': tokens.pop(0),
                })

        # finally we have our text
        command.update({
            'text': ' '.join(tokens),
        })

        return command<commit_msg>define that var<commit_after>import os

from django.http import HttpResponse, HttpResponseForbidden, JsonResponse
from django.views.decorators.csrf import csrf_exempt
from django.views.generic import View

from .models import Memory

SLACK_TOKEN = os.environ['SLACK_TOKEN']
BOT_NAME = 'membot'
KNOWN_COMMANDS = ['show',]

def homepage(request):
    return HttpResponse('Hello world this is membot')

class CommandView(View):
    response = {'text': ''}

    @csrf_exempt
    def dispatch(self, *args, **kwargs):
        return super(CommandView, self).dispatch(*args, **kwargs)

    def post(self, request, *args, **kwargs):
        received = request.POST
        auth_token = received.get('token', None)

        # make sure command is coming from the right place
        if auth_token != SLACK_TOKEN:
            return HttpResponseForbidden()

        # collect our command params
        command = self.parse_command(received)

        # if we don't actually have a command ...
        if not command['text']:
            self.set_response('Yes <@%s>?' % (command['person']))

        # otherwise take action
        else:
            if 'special' in command:
                action = command['special']

                if action not in KNOWN_COMMANDS:
                    self.set_response('\'%s\' sounds like a special command, <@%s>, but I haven\'t learned that one yet :(' % (command['special'], command['person']))

                if action == 'show':
                    memories = []
                    #TODO: handle `since` option

                    memories = Memory.objects.filter(is_active=True, category__in=command['categories']).distinct()

                    if memories:
                        intro = 'Here\'s what I remember about %s:\n' % ' '.join(command['categories'])
                        report = []

                        for memory in memories:
                            report.append('- On %s, <@%s> said: %s' % (memory.created.strftime('%B %d'), memory.person, memory.text))

                        self.set_response(intro + '```' + '\n'.join(report) + '```')
                    else:
                        self.set_response('Sorry, I don\'t remember anything like that!')

            # we have a memory to log, so do it for each defined category
            else:
                for category in command['categories']:
                    kwargs = {
                        'text': command['text'],
                        'person': command['person'],
                        'category': category,
                    }
                    memory = Memory(**kwargs)
                    memory.save()

                self.set_response('Got it, <@%s>!' % (command['person']))

        return JsonResponse(self.response)

    def set_response(self, text):
        self.response['text'] = text

    def parse_command(self, received):
        # split the command text for cleaning
        tokens = received.get('text', '').split(' ')

        # we don't need the bot name, and using `startswith`
        # catches most natural language punctuation
        if tokens[0].lower().startswith(BOT_NAME):
            del tokens[0]

        # just in case we had punctuation after the trigger name
        if tokens and tokens[0].lower() in [',', ':', '!', '?']:
            del tokens[0]

        # collect the hashtags
        categories = []
        for token in tokens:
            if token.startswith('#'):
                categories.append(token.lower())

        # give ourselves a default category if nothing else
        if not categories:
            categories = ['#general']

        # begin the command dict
        command = {
            'person': received.get('user_name', None),
            'categories': categories,
        }

        # see if we have a special command
        if tokens and tokens[0].lower() in ['please', 'plz']:
            del tokens[0]

            # take the next word for our special command
            if tokens:
                command.update({
                    'special': tokens.pop(0),
                })

        # finally we have our text
        command.update({
            'text': ' '.join(tokens),
        })

        return command<|endoftext|>