<commit_before># SNMP Simulator, http://snmpsim.sourceforge.net
# Managed value variation module
# Simulate a numeric value
# Valid values in module options are:
#   2  - Integer
#   65 - Counter32
#   66 - Gauge32
#   67 - TimeTicks
#   70 - Counter64
import sys
import math
import time
import random
from pysnmp.proto import rfc1902
from snmpsim import error

settingsCache = {}
valuesCache = {}
moduleOptions = {}
moduleContext = {}

booted = time.time()

def init(snmpEngine, **context):
    for k,v in [x.split(':') for x in context['options'].split(',')]:
        if k == 'addon':
            if k in moduleOptions:
                moduleOptions[k].append(v)
            else:
                moduleOptions[k] = [v]
        else:
            moduleOptions[k] = v
    if context['mode'] == 'recording':
        if 'iterations' in moduleOptions:
            moduleOptions['iterations'] = int(moduleOptions['iterations'])
            if moduleOptions['iterations']:
                moduleOptions['iterations'] = 1  # no reason for more
        if 'period' in moduleOptions:
            moduleOptions['period'] = float(moduleOptions['period'])
        else:
            moduleOptions['period'] = 10.0
    random.seed()

def variate(oid, tag, value, **context):
    if not context['nextFlag'] and not context['exactMatch']:
        return context['origOid'], tag, context['errorStatus']
    if context['setFlag']:
        return context['origOid'], tag, context['errorStatus']

    if oid not in settingsCache:
        settingsCache[oid] = dict([ x.split('=') for x in value.split(',') ])
        if 'min' not in settingsCache[oid]:
            settingsCache[oid]['min'] = 0
        if 'max' not in settingsCache[oid]:
            if tag == '67':
                settingsCache[oid]['max'] = 0xffffffffffffffff
            else:
                settingsCache[oid]['max'] = 0xffffffff

    if oid not in valuesCache:
        valuesCache[oid] = float(settingsCache[oid].get('initial', settingsCache[oid]['min'])), booted

    v, t = valuesCache[oid]

    if 'function' in settingsCache[oid]:
        f = getattr(math, settingsCache[oid]['function'])
    else:
        f = lambda x: x

    v += f((time.time() - t) * float(settingsCache[oid].get('rate', 1)))

    if 'increasing' in settingsCache[oid]:
        v = abs(v)

    v *= float(settingsCache[oid].get('scale', 1)) + float(settingsCache[oid].get('offset', 0))

    d = int(settingsCache[oid].get('deviation', 0))
    if d:
        if 'increasing' in settingsCache[oid]:
            v += random.randrange(0, d)
        else:
            v += random.randrange(-d, d)

    if v < settingsCache[oid]['min']:
        v = settingsCache[oid]['min']
    elif v > settingsCache[oid]['max']:
        if 'wrap' in settingsCache[oid]:
            v = settingsCache[oid]['min']
        else:
            v = settingsCache[oid]['max']

    if 'increasing' in settingsCache[oid]:
        valuesCache[oid] = v, time.time()

    return oid, tag, v

def record(oid, tag, value, **context):
    if 'started' not in moduleContext:
        moduleContext['started'] = time.time()
    if 'taglist' in moduleOptions and tag in moduleOptions['taglist']:
        if context['origValue'].tagSet in (rfc1902.Counter32.tagSet,
                                           rfc1902.Counter64.tagSet,
                                           rfc1902.TimeTicks.tagSet,
                                           rfc1902.Gauge32.tagSet,
                                           rfc1902.Integer.tagSet):
            tag += ':numeric'
            settings = { 'initial': value }
            if context['origValue'].tagSet not in (rfc1902.Gauge32.tagSet,
                                                   rfc1902.Integer.tagSet):
                settings['increasing'] = 1
            if context['origValue'].tagSet == rfc1902.TimeTicks.tagSet:
                settings['rate'] = 100
            if context['origValue'].tagSet == rfc1902.Integer.tagSet:
                settings['rate'] = 0
            if 'addon' in moduleOptions:
                settings.update(
                    dict([x.split('=') for x in moduleOptions['addon']])
                )

            value = ','.join([ '%s=%s' % (k,v) for k,v in settings.items() ])

            if 'hextag' in context:
                del context['hextag']
            if 'hexvalue' in context:
                del context['hexvalue']

            if 'iterations' in moduleOptions and oid not in moduleContext:
                moduleContext[oid] = settings

    if 'hextag' in context and context['hextag']:
        tag = context['hextag']
    if 'hexvalue' in context and context['hexvalue']:
        value = context['hexvalue']

    if 'iterations' in moduleOptions:
        if moduleOptions['iterations']:
            if context['stopFlag']:
                wait = max(0, float(moduleOptions['period']) - (time.time() - moduleContext['started']))
                while wait > 0:
                    sys.stdout.write('numeric: waiting %.2f sec(s), %s OIDs dumped, %s iterations remaining...\r' % (wait, context['total']+context['count'], moduleOptions['iterations']))
                    sys.stdout.flush()
                    time.sleep(1)
                    wait -= 1
                sys.stdout.write(' ' * 77 + '\r')
                moduleOptions['iterations'] -= 1
                moduleContext['started'] = time.time()
                raise error.MoreDataNotification()
            else:
                if oid in moduleContext:
                    moduleContext[oid]['time'] = time.time()
                    moduleContext[oid]['value'] = context['origValue']
                raise error.NoDataNotification()
        else:
            if oid in moduleContext:
                moduleContext[oid]['rate'] = \
                    (context['origValue'] - moduleContext[oid]['value'])/\
                    (time.time() - moduleContext[oid]['time'])

                del moduleContext[oid]['value'];
                del moduleContext[oid]['time'];

                value = ','.join(
                    [ '%s=%s' % (k,v) for k,v in moduleContext[oid].items() ]
                )

            return oid, tag, value
    else:
        return oid, tag, value

def shutdown(snmpEngine, **context): pass
<commit_msg>module options are optional<commit_after># SNMP Simulator, http://snmpsim.sourceforge.net
# Managed value variation module
# Simulate a numeric value
# Valid values in module options are:
#   2  - Integer
#   65 - Counter32
#   66 - Gauge32
#   67 - TimeTicks
#   70 - Counter64
import sys
import math
import time
import random
from pysnmp.proto import rfc1902
from snmpsim import error

settingsCache = {}
valuesCache = {}
moduleOptions = {}
moduleContext = {}

booted = time.time()

def init(snmpEngine, **context):
    if context['options']:
        for k,v in [x.split(':') for x in context['options'].split(',')]:
            if k == 'addon':
                if k in moduleOptions:
                    moduleOptions[k].append(v)
                else:
                    moduleOptions[k] = [v]
            else:
                moduleOptions[k] = v
    if context['mode'] == 'recording':
        if 'iterations' in moduleOptions:
            moduleOptions['iterations'] = int(moduleOptions['iterations'])
            if moduleOptions['iterations']:
                moduleOptions['iterations'] = 1  # no reason for more
        if 'period' in moduleOptions:
            moduleOptions['period'] = float(moduleOptions['period'])
        else:
            moduleOptions['period'] = 10.0
    random.seed()

def variate(oid, tag, value, **context):
    if not context['nextFlag'] and not context['exactMatch']:
        return context['origOid'], tag, context['errorStatus']
    if context['setFlag']:
        return context['origOid'], tag, context['errorStatus']

    if oid not in settingsCache:
        settingsCache[oid] = dict([ x.split('=') for x in value.split(',') ])
        if 'min' not in settingsCache[oid]:
            settingsCache[oid]['min'] = 0
        if 'max' not in settingsCache[oid]:
            if tag == '67':
                settingsCache[oid]['max'] = 0xffffffffffffffff
            else:
                settingsCache[oid]['max'] = 0xffffffff

    if oid not in valuesCache:
        valuesCache[oid] = float(settingsCache[oid].get('initial', settingsCache[oid]['min'])), booted

    v, t = valuesCache[oid]

    if 'function' in settingsCache[oid]:
        f = getattr(math, settingsCache[oid]['function'])
    else:
        f = lambda x: x

    v += f((time.time() - t) * float(settingsCache[oid].get('rate', 1)))

    if 'increasing' in settingsCache[oid]:
        v = abs(v)

    v *= float(settingsCache[oid].get('scale', 1)) + float(settingsCache[oid].get('offset', 0))

    d = int(settingsCache[oid].get('deviation', 0))
    if d:
        if 'increasing' in settingsCache[oid]:
            v += random.randrange(0, d)
        else:
            v += random.randrange(-d, d)

    if v < settingsCache[oid]['min']:
        v = settingsCache[oid]['min']
    elif v > settingsCache[oid]['max']:
        if 'wrap' in settingsCache[oid]:
            v = settingsCache[oid]['min']
        else:
            v = settingsCache[oid]['max']

    if 'increasing' in settingsCache[oid]:
        valuesCache[oid] = v, time.time()

    return oid, tag, v

def record(oid, tag, value, **context):
    if 'started' not in moduleContext:
        moduleContext['started'] = time.time()
    if 'taglist' in moduleOptions and tag in moduleOptions['taglist']:
        if context['origValue'].tagSet in (rfc1902.Counter32.tagSet,
                                           rfc1902.Counter64.tagSet,
                                           rfc1902.TimeTicks.tagSet,
                                           rfc1902.Gauge32.tagSet,
                                           rfc1902.Integer.tagSet):
            tag += ':numeric'
            settings = { 'initial': value }
            if context['origValue'].tagSet not in (rfc1902.Gauge32.tagSet,
                                                   rfc1902.Integer.tagSet):
                settings['increasing'] = 1
            if context['origValue'].tagSet == rfc1902.TimeTicks.tagSet:
                settings['rate'] = 100
            if context['origValue'].tagSet == rfc1902.Integer.tagSet:
                settings['rate'] = 0
            if 'addon' in moduleOptions:
                settings.update(
                    dict([x.split('=') for x in moduleOptions['addon']])
                )

            value = ','.join([ '%s=%s' % (k,v) for k,v in settings.items() ])

            if 'hextag' in context:
                del context['hextag']
            if 'hexvalue' in context:
                del context['hexvalue']

            if 'iterations' in moduleOptions and oid not in moduleContext:
                moduleContext[oid] = settings

    if 'hextag' in context and context['hextag']:
        tag = context['hextag']
    if 'hexvalue' in context and context['hexvalue']:
        value = context['hexvalue']

    if 'iterations' in moduleOptions:
        if moduleOptions['iterations']:
            if context['stopFlag']:
                wait = max(0, float(moduleOptions['period']) - (time.time() - moduleContext['started']))
                while wait > 0:
                    sys.stdout.write('numeric: waiting %.2f sec(s), %s OIDs dumped, %s iterations remaining...\r' % (wait, context['total']+context['count'], moduleOptions['iterations']))
                    sys.stdout.flush()
                    time.sleep(1)
                    wait -= 1
                sys.stdout.write(' ' * 77 + '\r')
                moduleOptions['iterations'] -= 1
                moduleContext['started'] = time.time()
                raise error.MoreDataNotification()
            else:
                if oid in moduleContext:
                    moduleContext[oid]['time'] = time.time()
                    moduleContext[oid]['value'] = context['origValue']
                raise error.NoDataNotification()
        else:
            if oid in moduleContext:
                moduleContext[oid]['rate'] = \
                    (context['origValue'] - moduleContext[oid]['value'])/\
                    (time.time() - moduleContext[oid]['time'])

                del moduleContext[oid]['value'];
                del moduleContext[oid]['time'];

                value = ','.join(
                    [ '%s=%s' % (k,v) for k,v in moduleContext[oid].items() ]
                )

            return oid, tag, value
    else:
        return oid, tag, value

def shutdown(snmpEngine, **context): pass
<|endoftext|>


<commit_before>"""
VerseBot for reddit
By Matthieu Grieger
database.py
Copyright (c) 2015 Matthieu Grieger (MIT License)
"""

import sqlite3

from config import DATABASE_PATH
from regex import find_already_quoted_verses

_conn = None
cur = None

def connect(logger):
    """ Connect to SQLite database. """

    global _conn
    global cur
    try:
        _conn = sqlite3.connect(DATABASE_PATH)
        cur = _conn.cursor()
    except sqlite3.Error:
        logger.critical("Connection to database failed. Exiting...")
        exit(1)

def update_book_stats(new_books, is_edit_or_delete=False):
    """ Updates book_stats table with recently quoted books.
    Alternatively, this function is also used to remove book counts
    that were added by a comment that has been edited/deleted. """

    for book in new_books.items():
        if is_edit_or_delete:
            cur.execute("UPDATE book_stats SET t_count = t_count - %d "
                        "WHERE book = '%s'" % (book[1], book[0]))
        else:
            cur.execute("UPDATE book_stats SET t_count = t_count + %d "
                        "WHERE book = '%s'" % (book[1], book[0]))
    _conn.commit()

def update_subreddit_stats(new_subreddits, is_edit_or_delete=False):
    """ Updates subreddit_stats table with subreddits that have recently used
    VerseBot. Alternatively, this function is also used to remove subreddit
    counts that were added by a comment that has been edited/deleted. """

    for subreddit in new_subreddits.items():
        if is_edit_or_delete:
            cur.execute("UPDATE subreddit_stats "
                        "SET t_count = t_count - %(num)s "
                        "WHERE sub = '%(subreddit)s';"
                        % {"subreddit": subreddit[0], "num": subreddit[1]})
            cur.execute("DELETE FROM subreddit_stats "
                        "WHERE t_count = 0;"
                        % {"subreddit": subreddit[0], "num": subreddit[1]})
        else:
            cur.execute("UPDATE subreddit_stats "
                        "SET t_count = t_count + %(num)s "
                        "WHERE sub = '%(subreddit)s';"
                        % {"subreddit": subreddit[0], "num": subreddit[1]})
            cur.execute("INSERT INTO subreddit_stats (sub, t_count) "
                        "SELECT '%(subreddit)s', %(num)s WHERE NOT EXISTS "
                        "(SELECT 1 FROM subreddit_stats "
                        "WHERE sub = '%(subreddit)s');"
                        % {"subreddit": subreddit[0], "num": subreddit[1]})
    _conn.commit()

def update_translation_stats(translations, is_edit_or_delete=False):
    """ Updates translation_stats table with recently used translations.
    Alternatively, this function is also used to remove translation counts
    that were added by a comment that has been edited/deleted. """

    for translation in translations.items():
        trans = translation[0]
        count = translation[1]
        if is_edit_or_delete:
            cur.execute("UPDATE translation_stats SET t_count = t_count - %d "
                        "WHERE trans = '%s'" % (count, trans))
        else:
            cur.execute("UPDATE translation_stats SET t_count = t_count + %d "
                        "WHERE trans = '%s'" % (count, trans))
    _conn.commit()

def prepare_translation_list_update():
    """ Prepares the translation_stats table for a translation list update.
    For all translations, available is set to false. The translations that are
    currently available will later be set to true. """

    cur.execute("UPDATE translation_stats SET available = 0;")
    _conn.commit()

def update_translation_list(translations):
    """ Updates translation_stats table with new translations that have been
    added. This query will also reset the available column for ALL
    translations to false, and then reset them to true individually if the
    translation exists in the local list of translations. This prevents users
    from trying to use translations within the database that are not
    officially supported anymore. """

    prepare_translation_list_update()
    for translation in translations:
        cur.execute(
            "UPDATE translation_stats SET available = 1 "
            "WHERE trans = '%s';" % translation.abbreviation)
        cur.execute(
            "INSERT INTO translation_stats "
            "(trans, name, lang, has_ot, has_nt, has_deut, available) "
            "SELECT '%(tran)s', '%(tname)s', '%(language)s', "
            "%(ot)s, %(nt)s, %(deut)s, 1 WHERE NOT EXISTS "
            "(SELECT 1 FROM translation_stats WHERE trans = '%(tran)s');"
            % {"tran": translation.abbreviation,
               "tname": translation.name.replace("'", "''"),
               "language": translation.language,
               "ot": int(translation.has_ot),
               "nt": int(translation.has_nt),
               "deut": int(translation.has_deut)})
    _conn.commit()

def update_user_translation(username, ot_trans, nt_trans, deut_trans):
    """ Updates user_translation table with new custom default translations
    specified by the user. """

    cur.execute(
        "UPDATE user_translations "
        "SET ot_default = '%(ot)s', nt_default = '%(nt)s', "
        "deut_default = '%(deut)s', last_used = datetime('now') "
        "WHERE username = '%(name)s';"
        % {"name": username, "ot": ot_trans,
           "nt": nt_trans, "deut": deut_trans})
    cur.execute(
        "INSERT INTO user_translations "
        "(username, ot_default, nt_default, deut_default) "
        "SELECT '%(name)s', '%(ot)s', '%(nt)s', '%(deut)s' "
        "WHERE NOT EXISTS (SELECT 1 FROM user_translations "
        "WHERE username = '%(name)s');"
        % {"name": username, "ot": ot_trans,
           "nt": nt_trans, "deut": deut_trans})
    _conn.commit()

def get_user_translation(username, bible_section):
    """ Retrieves the default translation for the user in a certain section
    of the Bible. """

    if bible_section == "Old Testament":
        section = "ot_default"
    elif bible_section == "New Testament":
        section = "nt_default"
    else:
        section = "deut_default"
    cur.execute("SELECT %s FROM user_translations WHERE username = '%s';"
                % (section, str(username)))
    try:
        translation = cur.fetchone()[0]
    except TypeError:
        translation = None
    cur.execute("UPDATE user_translations SET last_used = datetime('now') "
                "WHERE username = '%s'" % str(username))
    _conn.commit()
    return translation

def clean_user_translations():
    """ Removes user translation entries that haven't been used for 90 days
    or more to save on database space. """

    cur.execute("DELETE FROM user_translations "
                "WHERE last_used < datetime('now', '-90 days');")
    _conn.commit()

def update_subreddit_translation(subreddit, ot_trans, nt_trans, deut_trans):
    """ Updates subreddit_translation table with new custom default
    translations specified by a moderator of a subreddit. """

    cur.execute(
        "INSERT INTO subreddit_translations "
        "(sub, ot_default, nt_default, deut_default) "
        "SELECT '%(subreddit)s', '%(ot)s', '%(nt)s', '%(deut)s' "
        "WHERE NOT EXISTS "
        "(SELECT 1 FROM subreddit_translations WHERE sub = '%(subreddit)s');"
        % {"subreddit": subreddit.lower(), "ot": ot_trans,
           "nt": nt_trans, "deut": deut_trans})
    _conn.commit()

def get_subreddit_translation(subreddit, bible_section):
    """ Retrieves the default translation for the subreddit in a certain
    section of the Bible. """

    if bible_section == "Old Testament":
        section = "ot_default"
    elif bible_section == "New Testament":
        section = "nt_default"
    else:
        section = "deut_default"
    cur.execute("SELECT %s FROM subreddit_translations WHERE sub = '%s';"
                % (section, subreddit.lower()))
    try:
        trans = cur.fetchone()[0]
        return trans
    except TypeError:
        return None

def is_valid_translation(translation, testament):
    """ Checks the translations table for the supplied translation, and
    determines whether it is valid for the book that the user wants to quote.
    If the translation is not valid (either it is not available, or it doesn't
    contain the book), the translation will either default to the subreddit
    default or the global default. """

    if testament == "Old Testament":
        testament = "has_ot"
    elif testament == "New Testament":
        testament = "has_nt"
    else:
        testament = "has_deut"
    cur.execute("SELECT %s, available FROM translation_stats "
                "WHERE trans = '%s';" % (testament, translation))
    try:
        in_testament, is_available = cur.fetchone()
        if in_testament and bool(is_available):
            return True
        else:
            return False
    except TypeError:
        return False

def increment_comment_count():
    """ Increments the comment count entry whenever a new comment
    has been made. """

    cur.execute("UPDATE comment_count SET t_count = t_count + 1;")
    _conn.commit()

def decrement_comment_count():
    """ Decrements the comment count entry whenever a comment is deleted. """

    cur.execute("UPDATE comment_count SET t_count = t_count - 1;")
    _conn.commit()

def update_db_stats(verse_list):
    """ Iterates through all verses in verse_list and adds them to dicts
    to pass to the database update functions. """

    books = dict()
    translations = dict()
    subreddits = dict()

    for verse in verse_list:
        book = verse.book
        translation = verse.translation
        subreddit = verse.subreddit

        if book in books:
            books[book] += 1
        else:
            books[book] = 1

        if translation in translations:
            translations[translation] += 1
        else:
            translations[translation] = 1

        if subreddit in subreddits:
            subreddits[subreddit] += 1
        else:
            subreddits[subreddit] = 1

    update_book_stats(books)
    update_translation_stats(translations)
    update_subreddit_stats(subreddits)

def remove_invalid_statistics(message, subreddit):
    """ Performs necessary database operations to fix invalid statistics after
    a user has requested a comment to be edited or deleted. """

    invalid_verses = find_already_quoted_verses(message)
    invalid_books = dict()
    invalid_trans = dict()
    invalid_sub = dict()

    for verse in invalid_verses:
        book = verse[0].rstrip()
        translation = verse[1]
        if book in invalid_books:
            invalid_books[book] += 1
        else:
            invalid_books[book] = 1

        if translation in invalid_trans:
            invalid_trans[translation] += 1
        else:
            invalid_trans[translation] = 1

        if subreddit in invalid_sub:
            invalid_sub[subreddit] += 1
        else:
            invalid_sub[subreddit] = 1

    update_book_stats(invalid_books, is_edit_or_delete=True)
    update_translation_stats(invalid_trans, is_edit_or_delete=True)
    update_subreddit_stats(invalid_sub, is_edit_or_delete=True)
<commit_msg>Docstring additions and general formatting<commit_after>"""
VerseBot for reddit
By Matthieu Grieger
database.py
Copyright (c) 2015 Matthieu Grieger (MIT License)
"""

import sqlite3
from config import DATABASE_PATH
from regex import find_already_quoted_verses

_conn = None
cur = None

def connect(logger):
    """ Connect to SQLite database.

    :param logger: The logger used in the instance of VerseBot
    """
    global _conn
    global cur
    try:
        _conn = sqlite3.connect(DATABASE_PATH)
        cur = _conn.cursor()
    except sqlite3.Error:
        logger.critical("Connection to database failed. Exiting...")
        exit(1)

def update_book_stats(new_books, is_edit_or_delete=False):
    """ Updates book_stats table with recently quoted books.
    Alternatively, this function is also used to remove book counts
    that were added by a comment that has been edited/deleted.

    :param is_edit_or_delete: Whether a comment has been edited or deleted
    :param new_books: Recently quoted books
    """
    for book in new_books.items():
        if is_edit_or_delete:
            cur.execute("UPDATE book_stats SET t_count = t_count - %d "
                        "WHERE book = %s" % (book[1], book[0]))
        else:
            cur.execute("UPDATE book_stats SET t_count = t_count + %d "
                        "WHERE book = %s" % (book[1], book[0]))
    _conn.commit()

def update_subreddit_stats(new_subreddits, is_edit_or_delete=False):
    """ Updates subreddit_stats table with subreddits that have recently used
    VerseBot. Alternatively, this function is also used to remove subreddit
    counts that were added by a comment that has been edited/deleted.

    :param is_edit_or_delete: Whether a comment has been edited or deleted
    :param new_subreddits: Subreddits that have recently used VerseBot
    """
    for subreddit in new_subreddits.items():
        if is_edit_or_delete:
            cur.execute("UPDATE subreddit_stats SET t_count = t_count - %d "
                        "WHERE sub = %s;" % (subreddit[1], subreddit[0]))
            cur.execute("DELETE FROM subreddit_stats WHERE t_count = 0;")
        else:
            cur.execute("UPDATE subreddit_stats SET t_count = t_count + %d "
                        "WHERE sub = %s;" % (subreddit[1], subreddit[0]))
            cur.execute("INSERT INTO subreddit_stats (sub, t_count) "
                        "SELECT %(subreddit)s, %(num)s WHERE NOT EXISTS "
                        "(SELECT 1 FROM subreddit_stats "
                        "WHERE sub = %(subreddit)s);"
                        % {"subreddit": subreddit[0], "num": subreddit[1]})
    _conn.commit()

def update_translation_stats(translations, is_edit_or_delete=False):
    """ Updates translation_stats table with recently used translations.
    Alternatively, this function is also used to remove translation counts
    that were added by a comment that has been edited/deleted.

    :param is_edit_or_delete: Whether a comment has been edited or deleted
    :param translations: Recently used translations
    """
    for translation in translations.items():
        trans = translation[0]
        count = translation[1]
        if is_edit_or_delete:
            cur.execute("UPDATE translation_stats SET t_count = t_count - %d "
                        "WHERE trans = %s" % (count, trans))
        else:
            cur.execute("UPDATE translation_stats SET t_count = t_count + %d "
                        "WHERE trans = %s" % (count, trans))
    _conn.commit()

def prepare_translation_list_update():
    """ Prepares the translation_stats table for a translation list update.
    For all translations, available is set to false. The translations that are
    currently available will later be set to true. """

    cur.execute("UPDATE translation_stats SET available = 0;")
    _conn.commit()

def update_translation_list(translations):
    """ Updates translation_stats table with new translations that have been
    added. This query will also reset the available column for ALL
    translations to false, and then reset them to true individually if the
    translation exists in the local list of translations. This prevents users
    from trying to use translations within the database that are not
    officially supported anymore.

    :param translations: Available translations
    """
    prepare_translation_list_update()
    for translation in translations:
        cur.execute("UPDATE translation_stats SET available = 1 "
                    "WHERE trans = %s;" % translation.abbreviation)
        cur.execute("INSERT INTO translation_stats "
                    "(trans, name, lang, has_ot, has_nt, has_deut, available) "
                    "SELECT %(tran)s, %(tname)s, %(language)s, "
                    "%(ot)s, %(nt)s, %(deut)s, 1 WHERE NOT EXISTS "
                    "(SELECT 1 FROM translation_stats WHERE trans = %(tran)s);"
                    % {"tran": translation.abbreviation,
                       "tname": translation.name.replace("'", "''"),
                       "language": translation.language,
                       "ot": int(translation.has_ot),
                       "nt": int(translation.has_nt),
                       "deut": int(translation.has_deut)})
    _conn.commit()

def update_user_translation(username, ot_trans, nt_trans, deut_trans):
    """ Updates user_translation table with new custom default translations
    specified by the user.

    :param deut_trans: Requested Deuterocanon translation
    :param nt_trans: Requested New Testament translation
    :param ot_trans: Requested Old Testament translation
    :param username: Username of the user who requested the update
    """
    cur.execute("UPDATE user_translations "
                "SET ot_default = %s, nt_default = %s, "
                "deut_default = %s, last_used = datetime('now') "
                "WHERE username = %s;"
                % (ot_trans, nt_trans, deut_trans, username))
    cur.execute("INSERT INTO user_translations "
                "(username, ot_default, nt_default, deut_default) "
                "SELECT %(name)s, %(ot)s, %(nt)s, %(deut)s "
                "WHERE NOT EXISTS (SELECT 1 FROM user_translations "
                "WHERE username = %(name)s);"
                % {"name": username, "ot": ot_trans,
                   "nt": nt_trans, "deut": deut_trans})
    _conn.commit()

def get_user_translation(username, bible_section):
    """ Retrieves the default translation for the user in a certain section
    of the Bible.

    :param bible_section: Old Testament, New Testament, or Deuterocanon
    :param username: Username of the user who called VerseBot
    """
    if bible_section == "Old Testament":
        section = "ot_default"
    elif bible_section == "New Testament":
        section = "nt_default"
    else:
        section = "deut_default"
    cur.execute("SELECT %s FROM user_translations WHERE username = %s;"
                % (section, str(username)))
    try:
        translation = cur.fetchone()[0]
    except TypeError:
        translation = None
    cur.execute("UPDATE user_translations SET last_used = datetime('now') "
                "WHERE username = %s" % str(username))
    _conn.commit()
    return translation

def clean_user_translations():
    """ Removes user translation entries that haven't been used for 90 days
    or more to save on database space. """

    cur.execute("DELETE FROM user_translations "
                "WHERE last_used < datetime('now', '-90 days');")
    _conn.commit()

def update_subreddit_translation(subreddit, ot_trans, nt_trans, deut_trans):
    """ Updates subreddit_translation table with new custom default
    translations specified by a moderator of a subreddit.

    :param deut_trans: Requested Deuterocanon translation
    :param nt_trans: Requested New Testament translation
    :param ot_trans: Requested Old Testament translation
    :param subreddit: Subreddit that requested a translation update
    """
    cur.execute("INSERT INTO subreddit_translations "
                "(sub, ot_default, nt_default, deut_default) "
                "SELECT %(subreddit)s, %(ot)s, %(nt)s, %(deut)s "
                "WHERE NOT EXISTS (SELECT 1 FROM subreddit_translations "
                "WHERE sub = %(subreddit)s);"
                % {"subreddit": subreddit.lower(), "ot": ot_trans,
                   "nt": nt_trans, "deut": deut_trans})
    _conn.commit()

def get_subreddit_translation(subreddit, bible_section):
    """ Retrieves the default translation for the subreddit in a certain
    section of the Bible.

    :param bible_section: Old Testament, New Testament, or Deuterocanon
    :param subreddit: The subreddit that called VerseBot
    """
    if bible_section == "Old Testament":
        section = "ot_default"
    elif bible_section == "New Testament":
        section = "nt_default"
    else:
        section = "deut_default"
    cur.execute("SELECT %s FROM subreddit_translations WHERE sub = %s;"
                % (section, subreddit.lower()))
    try:
        trans = cur.fetchone()[0]
        return trans
    except TypeError:
        return None

def is_valid_translation(translation, testament):
    """ Checks the translations table for the supplied translation, and
    determines whether it is valid for the book that the user wants to quote.
    If the translation is not valid (either it is not available, or it doesn't
    contain the book), the translation will either default to the subreddit
    default or the global default.

    :param testament: Old Testament, New Testament, or Deuterocanon
    :param translation: Requested translation
    """
    if testament == "Old Testament":
        testament = "has_ot"
    elif testament == "New Testament":
        testament = "has_nt"
    else:
        testament = "has_deut"
    cur.execute("SELECT %s, available FROM translation_stats "
                "WHERE trans = %s;" % (testament, translation))
    try:
        in_testament, is_available = cur.fetchone()
        if in_testament and bool(is_available):
            return True
        else:
            return False
    except TypeError:
        return False

def increment_comment_count():
    """ Increments the comment count entry whenever a new comment
    has been made. """

    cur.execute("UPDATE comment_count SET t_count = t_count + 1;")
    _conn.commit()

def decrement_comment_count():
    """ Decrements the comment count entry whenever a comment is deleted. """

    cur.execute("UPDATE comment_count SET t_count = t_count - 1;")
    _conn.commit()

def update_db_stats(verse_list):
    """ Iterates through all verses in verse_list and adds them to dicts
    to pass to the database update functions.

    :param verse_list: List of recently quoted verses
    """
    books = dict()
    translations = dict()
    subreddits = dict()

    for verse in verse_list:
        book = verse.book
        translation = verse.translation
        subreddit = verse.subreddit

        if book in books:
            books[book] += 1
        else:
            books[book] = 1

        if translation in translations:
            translations[translation] += 1
        else:
            translations[translation] = 1

        if subreddit in subreddits:
            subreddits[subreddit] += 1
        else:
            subreddits[subreddit] = 1

    update_book_stats(books)
    update_translation_stats(translations)
    update_subreddit_stats(subreddits)

def remove_invalid_statistics(message, subreddit):
    """ Performs necessary database operations to fix invalid statistics after
    a user has requested a comment to be edited or deleted.

    :param subreddit: Subreddit where the quotation is located
    :param message: Message containing verse quotations
    """
    invalid_verses = find_already_quoted_verses(message)
    invalid_books = dict()
    invalid_trans = dict()
    invalid_sub = dict()

    for verse in invalid_verses:
        book = verse[0].rstrip()
        translation = verse[1]
        if book in invalid_books:
            invalid_books[book] += 1
        else:
            invalid_books[book] = 1

        if translation in invalid_trans:
            invalid_trans[translation] += 1
        else:
            invalid_trans[translation] = 1

        if subreddit in invalid_sub:
            invalid_sub[subreddit] += 1
        else:
            invalid_sub[subreddit] = 1

    update_book_stats(invalid_books, is_edit_or_delete=True)
    update_translation_stats(invalid_trans, is_edit_or_delete=True)
    update_subreddit_stats(invalid_sub, is_edit_or_delete=True)
<|endoftext|>


<commit_before>#===================================================================================================
#  C L A S S
#===================================================================================================
import os, subprocess, re, signal, sys, MySQLdb, json
import datetime, time
import phedexDataset
import phedexApi

class Alarm(Exception):
    pass

def alarm_handler(signum, frame):
    raise Alarm

class PhedexDataHandler:
    def __init__(self,allSites):
        self.newAccess = False
        self.phedexGroup =  os.environ['DETOX_GROUP']
        self.phedexDatasets = {}
        self.otherDatasets = {}
        self.runAwayGroups = {}
        self.allSites = allSites
        self.epochTime = int(time.time())

    def shouldAccessPhedex(self):
        # number of hours until it will rerun
        renewMinInterval = int(os.environ.get('DETOX_CYCLE_HOURS'))
        statusDir = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS']
        fileName = statusDir + '/' + os.environ['DETOX_PHEDEX_CACHE']
        if not os.path.isfile(fileName):
            return True
        if not os.path.getsize(fileName) > 0:
            return True

        timeNow = datetime.datetime.now()
        deltaNhours = datetime.timedelta(seconds = 60*60*(renewMinInterval-1))
        modTime = datetime.datetime.fromtimestamp(os.path.getmtime(fileName))
        if (timeNow-deltaNhours) < modTime:
            print "  -- last time cache renewed on " + str(modTime)
            return False
        return True

    def extractPhedexData(self,federation):
        if self.shouldAccessPhedex() :
            try:
                self.retrievePhedexData(federation)
                self.newAccess = True
            except:
                raise
        else:
            print "  -- reading from cache --"

        self.readPhedexData()

    def retrievePhedexData(self,federation):
        phedexDatasets = {}
        webServer = 'https://cmsweb.cern.ch/'
        phedexBlocks = 'phedex/datasvc/json/prod/blockreplicas'
        args = 'show_dataset=y&subscribed=y&' + federation

        url = '"'+webServer+phedexBlocks+'?'+args+'"'
        cmd = 'curl -k -H "Accept: text/xml" ' + url

        print ' Access phedexDb: ' + cmd
        tmpname = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/tmp.txt'
        tmpfile = open(tmpname, "w")

        process = subprocess.Popen(cmd, stdout=tmpfile, stderr=subprocess.PIPE,
                                   bufsize=4096,shell=True)

        signal.signal(signal.SIGALRM, alarm_handler)
        signal.alarm(30*60)  # 30 minutes
        try:
            strout, error = process.communicate()
            tmpfile.close()
            signal.alarm(0)
        except Alarm:
            print " Oops, taking too long!"
            raise Exception(" FATAL -- Call to PhEDEx timed out, stopping")

        if process.returncode != 0:
            print " Received non-zero exit status: " + str(process.returncode)
            raise Exception(" FATAL -- Call to PhEDEx failed, stopping")

        tmpfile = open(tmpname, "r")
        strout = tmpfile.readline()
        tmpfile.close()
        #os.remove(tmpname)
        dataJson = json.loads(strout)
        datasets = (dataJson["phedex"])["dataset"]
        for dset in datasets:
            datasetName = dset["name"]

            user = re.findall(r"USER",datasetName)
            blocks = dset["block"]
            for block in blocks:
                replicas = block["replica"]
                for siterpl in replicas:
                    group = siterpl["group"]
                    if group == 'IB RelVal':
                        group = 'IB-RelVal'

                    site = str(siterpl["node"])
                    #if site not in self.allSites:
                    #    continue

                    if datasetName not in phedexDatasets:
                        phedexDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                    dataset = phedexDatasets[datasetName]

                    if site == 'T2_US_CALTECH':
                        if datasetName == '/QCD_Pt-10to20_EMEnriched_Tune4C_13TeV_pythia8/Spring14miniaod-castor_PU20bx25_POSTLS170_V5-v1/MINIAODSIM':
                            print siterpl
                    size = float(siterpl["bytes"])/1000/1000/1000
                    strdone = siterpl["complete"]
                    isdone = 1
                    if strdone == 'n' :
                        isdone = 0
                    cust = siterpl["custodial"]
                    reqtime = int(float(siterpl["time_create"]))
                    updtime = int(float(siterpl["time_update"]))
                    files = int(siterpl["files"])

                    iscust = 0
                    if len(user) > 0 or cust == 'y':
                        iscust = 1
                    dataset.updateForSite(site,size,group,files,iscust,reqtime,updtime,isdone)

        # Create our local cache files of the status per site
        filename = os.environ['DETOX_PHEDEX_CACHE']
        outputFile = open(os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/'
                          + filename, "w")
        for datasetName in phedexDatasets:
            phedexDatasets[datasetName].setFinalValues()
            line = phedexDatasets[datasetName].printIntoLine()
            #any dataset line should be above 10 characters
            if len(line) < 10:
                print " SKIPING " + datasetName
                continue
            outputFile.write(line)
        outputFile.close()

    def readPhedexData(self):
        filename = os.environ['DETOX_PHEDEX_CACHE']
        fileName = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/' + filename
	if not os.path.exists(fileName):
	    return phedexDatasets

	inputFile = open(fileName,'r')
        for line in inputFile.xreadlines():
            items = line.split()
            datasetName = items[0]
            group = items[1]
            siteName = items[5]
            size = float(items[2])

            if self.allSites[siteName].getStatus() == 0:
                continue
            if group == 'local':
                continue

            dataset = None
            if group != self.phedexGroup:
                if datasetName not in self.otherDatasets:
                    self.otherDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                dataset = self.otherDatasets[datasetName]
            else:
                if datasetName not in self.phedexDatasets:
                    self.phedexDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                dataset = self.phedexDatasets[datasetName]
            dataset.fillFromLine(line)
            dataset.setFinalValues()
        inputFile.close()

        if self.phedexGroup == 'AnalysisOps':
            self.printRunawaySets()

    def getPhedexDatasets(self):
        return self.phedexDatasets

    def getPhedexDatasetsAtSite(self,site):
        dsets = []
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets.append(self.phedexDatasets[datasetName])
        return dsets

    def getDatasetsAtSite(self,site):
        dsets = []
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets.append(datasetName)
        return dsets

    def getDatasetsByRank(self,site):
        dsets = {}
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets[datasetName] = dataset.getLocalRank(site)
        return sorted(dsets,key=dsets.get,reverse=True)

    def renewedCache(self):
        return self.newAccess

    def getRunAwayGroups(self,site):
        groups = []
        if site in self.runAwayGroups:
            groups = self.runAwayGroups[site]
        return groups

    def getRunAwaySets(self,site):
        runAway = {}
        for dset in self.phedexDatasets:
            if dset not in self.otherDatasets:
                continue

            dataset = self.otherDatasets[dset]
            if not dataset.isOnSite(site):
                continue
            group = dataset.group(site)
            runAway[dset] = group
        return runAway

    def printRunawaySets(self):
        siteSizes = {}
        siteSets = {}
        for dset in self.phedexDatasets:
            if dset not in self.otherDatasets:
                continue

            dataset = self.otherDatasets[dset]
            for site in dataset.siteNames:
                group = dataset.group(site)
                size  = dataset.size(site)
                if site not in siteSizes:
                    siteSizes[site] = 0
                    siteSets[site] = 0
                siteSizes[site] = siteSizes[site] + size/1000
                siteSets[site] = siteSets[site] + 1
                if site not in self.runAwayGroups:
                    dsetName = dataset.dataset
                    if 'AOD' in dsetName:
#                        if group != 'RelVal' and group != 'DataOps':
                         if not site.startswith('T1_'):
                            if group == 'local':
                               continue
                            print site
                            print dataset.dataset
                            print group + '--> AnalysisOps'
                            phedex = phedexApi.phedexApi(logPath='./')
                            check,response = phedex.changeGroup(site,dsetName,'AnalysisOps')
                            if check:
                                print " ERROR - phedexApi.updateRequest failed"
                                print response
                                del phedex
                            else:
                                del phedex
                                continue

                    self.runAwayGroups[site] = [group]
                else:
                    if group not in self.runAwayGroups[site]:
                        self.runAwayGroups[site].append(group)

        if (len(siteSizes) < 1):
            return
        print " !! WARNING !! - those sites have datasets in wrong groups"
        for site in sorted(siteSizes):
            print ' %3d %6.2f TB'%(siteSets[site],siteSizes[site]) + ": " + site

<commit_msg>fix bug<commit_after>#===================================================================================================
#  C L A S S
#===================================================================================================
import os, subprocess, re, signal, sys, MySQLdb, json
import datetime, time
import phedexDataset
import phedexApi

class Alarm(Exception):
    pass

def alarm_handler(signum, frame):
    raise Alarm

class PhedexDataHandler:
    def __init__(self,allSites):
        self.newAccess = False
        self.phedexGroup =  os.environ['DETOX_GROUP']
        self.phedexDatasets = {}
        self.otherDatasets = {}
        self.runAwayGroups = {}
        self.allSites = allSites
        self.epochTime = int(time.time())

    def shouldAccessPhedex(self):
        # number of hours until it will rerun
        renewMinInterval = int(os.environ.get('DETOX_CYCLE_HOURS'))
        statusDir = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS']
        fileName = statusDir + '/' + os.environ['DETOX_PHEDEX_CACHE']
        if not os.path.isfile(fileName):
            return True
        if not os.path.getsize(fileName) > 0:
            return True

        timeNow = datetime.datetime.now()
        deltaNhours = datetime.timedelta(seconds = 60*60*(renewMinInterval-1))
        modTime = datetime.datetime.fromtimestamp(os.path.getmtime(fileName))
        if (timeNow-deltaNhours) < modTime:
            print "  -- last time cache renewed on " + str(modTime)
            return False
        return True

    def extractPhedexData(self,federation):
        if self.shouldAccessPhedex() :
            try:
                self.retrievePhedexData(federation)
                self.newAccess = True
            except:
                raise
        else:
            print "  -- reading from cache --"

        self.readPhedexData()

    def retrievePhedexData(self,federation):
        phedexDatasets = {}
        webServer = 'https://cmsweb.cern.ch/'
        phedexBlocks = 'phedex/datasvc/json/prod/blockreplicas'
        args = 'show_dataset=y&subscribed=y&' + federation

        url = '"'+webServer+phedexBlocks+'?'+args+'"'
        cmd = 'curl -k -H "Accept: text/xml" ' + url

        print ' Access phedexDb: ' + cmd
        tmpname = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/tmp.txt'
        tmpfile = open(tmpname, "w")

        process = subprocess.Popen(cmd, stdout=tmpfile, stderr=subprocess.PIPE,
                                   bufsize=4096,shell=True)

        signal.signal(signal.SIGALRM, alarm_handler)
        signal.alarm(30*60)  # 30 minutes
        try:
            strout, error = process.communicate()
            tmpfile.close()
            signal.alarm(0)
        except Alarm:
            print " Oops, taking too long!"
            raise Exception(" FATAL -- Call to PhEDEx timed out, stopping")

        if process.returncode != 0:
            print " Received non-zero exit status: " + str(process.returncode)
            raise Exception(" FATAL -- Call to PhEDEx failed, stopping")

        tmpfile = open(tmpname, "r")
        strout = tmpfile.readline()
        tmpfile.close()
        #os.remove(tmpname)
        dataJson = json.loads(strout)
        datasets = (dataJson["phedex"])["dataset"]
        for dset in datasets:
            datasetName = dset["name"]

            user = re.findall(r"USER",datasetName)
            blocks = dset["block"]
            for block in blocks:
                replicas = block["replica"]
                for siterpl in replicas:
                    group = siterpl["group"]
                    if group == 'IB RelVal':
                        group = 'IB-RelVal'

                    site = str(siterpl["node"])
                    #if site not in self.allSites:
                    #    continue

                    if datasetName not in phedexDatasets:
                        phedexDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                    dataset = phedexDatasets[datasetName]

                    if site == 'T2_US_CALTECH':
                        if datasetName == '/QCD_Pt-10to20_EMEnriched_Tune4C_13TeV_pythia8/Spring14miniaod-castor_PU20bx25_POSTLS170_V5-v1/MINIAODSIM':
                            print siterpl
                    size = float(siterpl["bytes"])/1000/1000/1000
                    strdone = siterpl["complete"]
                    isdone = 1
                    if strdone == 'n' :
                        isdone = 0
                    cust = siterpl["custodial"]
                    reqtime = int(float(siterpl["time_create"]))
                    updtime = int(float(siterpl["time_update"]))
                    files = int(siterpl["files"])

                    iscust = 0
                    if len(user) > 0 or cust == 'y':
                        iscust = 1
                    dataset.updateForSite(site,size,group,files,iscust,reqtime,updtime,isdone)

        # Create our local cache files of the status per site
        filename = os.environ['DETOX_PHEDEX_CACHE']
        outputFile = open(os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/'
                          + filename, "w")
        for datasetName in phedexDatasets:
            phedexDatasets[datasetName].setFinalValues()
            line = phedexDatasets[datasetName].printIntoLine()
            #any dataset line should be above 10 characters
            if len(line) < 10:
                print " SKIPING " + datasetName
                continue
            outputFile.write(line)
        outputFile.close()

    def readPhedexData(self):
        filename = os.environ['DETOX_PHEDEX_CACHE']
        fileName = os.environ['DETOX_DB'] + '/' + os.environ['DETOX_STATUS'] + '/' + filename
	if not os.path.exists(fileName):
	    return phedexDatasets

	inputFile = open(fileName,'r')
        for line in inputFile.xreadlines():
            items = line.split()
            datasetName = items[0]
            group = items[1]
            siteName = items[5]
            size = float(items[2])

	    if siteName not in self.allSites:
                continue
            if self.allSites[siteName].getStatus() == 0:
                continue
            if group == 'local':
                continue

            dataset = None
            if group != self.phedexGroup:
                if datasetName not in self.otherDatasets:
                    self.otherDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                dataset = self.otherDatasets[datasetName]
            else:
                if datasetName not in self.phedexDatasets:
                    self.phedexDatasets[datasetName] = phedexDataset.PhedexDataset(datasetName)
                dataset = self.phedexDatasets[datasetName]
            dataset.fillFromLine(line)
            dataset.setFinalValues()
        inputFile.close()

        if self.phedexGroup == 'AnalysisOps':
            self.printRunawaySets()

    def getPhedexDatasets(self):
        return self.phedexDatasets

    def getPhedexDatasetsAtSite(self,site):
        dsets = []
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets.append(self.phedexDatasets[datasetName])
        return dsets

    def getDatasetsAtSite(self,site):
        dsets = []
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets.append(datasetName)
        return dsets

    def getDatasetsByRank(self,site):
        dsets = {}
        for datasetName in self.phedexDatasets.keys():
            dataset = self.phedexDatasets[datasetName]
            if dataset.isOnSite(site):
                dsets[datasetName] = dataset.getLocalRank(site)
        return sorted(dsets,key=dsets.get,reverse=True)

    def renewedCache(self):
        return self.newAccess

    def getRunAwayGroups(self,site):
        groups = []
        if site in self.runAwayGroups:
            groups = self.runAwayGroups[site]
        return groups

    def getRunAwaySets(self,site):
        runAway = {}
        for dset in self.phedexDatasets:
            if dset not in self.otherDatasets:
                continue

            dataset = self.otherDatasets[dset]
            if not dataset.isOnSite(site):
                continue
            group = dataset.group(site)
            runAway[dset] = group
        return runAway

    def printRunawaySets(self):
        siteSizes = {}
        siteSets = {}
        for dset in self.phedexDatasets:
            if dset not in self.otherDatasets:
                continue

            dataset = self.otherDatasets[dset]
            for site in dataset.siteNames:
                group = dataset.group(site)
                size  = dataset.size(site)
                if site not in siteSizes:
                    siteSizes[site] = 0
                    siteSets[site] = 0
                siteSizes[site] = siteSizes[site] + size/1000
                siteSets[site] = siteSets[site] + 1
                if site not in self.runAwayGroups:
                    dsetName = dataset.dataset
                    if 'AOD' in dsetName:
#                        if group != 'RelVal' and group != 'DataOps':
                         if not site.startswith('T1_'):
                            if group == 'local':
                               continue
                            print site
                            print dataset.dataset
                            print group + '--> AnalysisOps'
                            phedex = phedexApi.phedexApi(logPath='./')
                            check,response = phedex.changeGroup(site,dsetName,'AnalysisOps')
                            if check:
                                print " ERROR - phedexApi.updateRequest failed"
                                print response
                                del phedex
                            else:
                                del phedex
                                continue

                    self.runAwayGroups[site] = [group]
                else:
                    if group not in self.runAwayGroups[site]:
                        self.runAwayGroups[site].append(group)

        if (len(siteSizes) < 1):
            return
        print " !! WARNING !! - those sites have datasets in wrong groups"
        for site in sorted(siteSizes):
            print ' %3d %6.2f TB'%(siteSets[site],siteSizes[site]) + ": " + site

<|endoftext|>


<commit_before>class Request(object):
    """Request class."""

    def __init__(self, url, action, request, **params):
        self._action = action
        self._request = request
        self._params = params
        self._params['url'] = url

    def action(self):
        return self._action

    def request(self):
        return self._request

    def params(self):
        return self._params

class Response(object):
    """Response class."""

    def __init__(self, response):
        self._id = response.pop('requestId')
        self._action = response.pop('action')
        self._success = response.pop('status') == 'success'
        self._code = response.pop('code', None)
        self._error = response.pop('error', None)
        self._response = response

    def id(self):
        return self._id

    def action(self):
        return self._action

    def success(self):
        return self._success

    def code(self):
        return self._code

    def error(self):
        return self._error

    def response(self, key=None):
        if key:
            return self._response[key]
        return self._response

class ApiObject(object):
    """Api object class."""

    def __init__(self, transport):
        self._transport = transport
        self._transport_name = self._transport.name()

    def _http_transport(self):
        return self._transport_name == 'http'

    def _websocket_transport(self):
        return self._transport_name == 'websocket'

    def _request(self, url, action, request, **params):
        request = Request(url, action, request, **params)
        response = self._transport.request(request.action(), request.request(),
                                           **request.params())
        return Response(response)

    def _ensure_success_response(self, response, message):
        if not response.success():
            raise ResponseApiObjectException(message, self._transport_name,
                                             response.code(), response.error())

    def _ensure_http_transport(self):
        if not self._http_transport():
            raise ApiObjectException('Implemented only for http transport')

class ApiObjectException(Exception):
    """Api object exception."""
    pass

class ResponseApiObjectException(ApiObjectException):
    """Response api object exception."""

    def __init__(self, message, transport_name, code, error):
        message = '%s. Transport: %s. Code: %s. Error: %s.' % (message,
                                                               transport_name,
                                                               code, error)
        ApiObjectException.__init__(self, message)
<commit_msg>Rename ResponseApiObjectExcaption to ApiObjectResponseException<commit_after>class Request(object):
    """Request class."""

    def __init__(self, url, action, request, **params):
        self._action = action
        self._request = request
        self._params = params
        self._params['url'] = url

    def action(self):
        return self._action

    def request(self):
        return self._request

    def params(self):
        return self._params

class Response(object):
    """Response class."""

    def __init__(self, response):
        self._id = response.pop('requestId')
        self._action = response.pop('action')
        self._success = response.pop('status') == 'success'
        self._code = response.pop('code', None)
        self._error = response.pop('error', None)
        self._response = response

    def id(self):
        return self._id

    def action(self):
        return self._action

    def success(self):
        return self._success

    def code(self):
        return self._code

    def error(self):
        return self._error

    def response(self, key=None):
        if key:
            return self._response[key]
        return self._response

class ApiObject(object):
    """Api object class."""

    def __init__(self, transport):
        self._transport = transport
        self._transport_name = self._transport.name()

    def _http_transport(self):
        return self._transport_name == 'http'

    def _websocket_transport(self):
        return self._transport_name == 'websocket'

    def _request(self, url, action, request, **params):
        request = Request(url, action, request, **params)
        response = self._transport.request(request.action(), request.request(),
                                           **request.params())
        return Response(response)

    def _ensure_success_response(self, response, message):
        if not response.success():
            raise ApiObjectResponseException(message, self._transport_name,
                                             response.code(), response.error())

    def _ensure_http_transport(self):
        if not self._http_transport():
            raise ApiObjectException('Implemented only for http transport')

class ApiObjectException(Exception):
    """Api object exception."""
    pass

class ApiObjectResponseException(ApiObjectException):
    """Api object response exception."""

    def __init__(self, message, transport_name, code, error):
        message = '%s. Transport: %s. Code: %s. Error: %s.' % (message,
                                                               transport_name,
                                                               code, error)
        ApiObjectException.__init__(self, message)
<|endoftext|>


<commit_before>import json
import os

import amazon.api

from . import models

def file_credentials(path="~/.caudexer.aws.json"):
    full_path = os.path.expanduser(path)
    with open(full_path, 'r') as fp:
        return json.loads(fp.read())

def search(keywords):
    """
    :returns: *AmazonBook* models that match keywords
    """
    return [models.AmazonBook.from_product(result)
            for result in search_books(Keywords=keywords)]

def search_books(get_credentials=file_credentials, **options):
    """
    Yields book products that match search options
    """
    creds = get_credentials()
    api = amazon.api.AmazonAPI(**creds)
    options.setdefault('SearchIndex', 'Books')
    results = api.search(**options)
    for index, result_item in enumerate(results):
        yield result_item

<commit_msg>No amazon for you.<commit_after>import json
import os

import amazon.api

from . import models

def file_credentials(path="~/.caudexer.aws.json"):
    full_path = os.path.expanduser(path)
    with open(full_path, 'r') as fp:
        return json.loads(fp.read())

def search(keywords):
    """
    :returns: *AmazonBook* models that match keywords
    """
    return [models.AmazonBook.from_product(result)
            for result in search_books(Keywords=keywords)]

def search_books(get_credentials=file_credentials, **options):
    """
    Yields book products that match search options
    """
    try:
        creds = get_credentials()
        api = amazon.api.AmazonAPI(**creds)
    except Exception as exc:
        print(exc)
        return []
    options.setdefault('SearchIndex', 'Books')
    results = api.search(**options)
    for index, result_item in enumerate(results):
        yield result_item

<|endoftext|>


<commit_before>import os, sys
import vtk
from vtk.util.colors import tomato

this_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(this_dir + '/../..')

from tessagon.adaptors.vtk_adaptor import VtkAdaptor

from tessagon.types import *
from tessagon.misc.shapes import *

def main():
  ren = vtk.vtkRenderer()
  renWin = vtk.vtkRenderWindow()
  renWin.AddRenderer(ren)
  iren = vtk.vtkRenderWindowInteractor()
  iren.SetRenderWindow(renWin)

  ren.AddActor(hex_tessagon())
  ren.AddActor(tri_tessagon())
  ren.AddActor(rhombus_tessagon())
  ren.AddActor(octo_tessagon())
  ren.AddActor(hex_tri_tessagon())
  ren.AddActor(hex_square_tri_tessagon())
  ren.AddActor(square_tessagon())
  ren.AddActor(pythagorean_tessagon())
  ren.AddActor(brick_tessagon())
  ren.AddActor(dodeca_tessagon())

  ren.SetBackground(0.3, 0.3, 0.3)
  renWin.SetSize(800, 600)
  iren.Initialize()
  ren.ResetCamera()
  renWin.Render()
  iren.Start()

def tessellate(f, tessagon_class, **kwargs):
  extra_args = {'adaptor_class' : VtkAdaptor}
  tessagon = tessagon_class(f, **{**kwargs, **extra_args})

  poly_data = tessagon.create_mesh()
  mapper = vtk.vtkPolyDataMapper()
  mapper.SetInputData(poly_data)
  actor = vtk.vtkActor()
  actor.SetMapper(mapper)
  actor.GetProperty().SetColor(tomato)
  actor.GetProperty().SetEdgeColor(0.0, 0.0, 0.0)
  actor.GetProperty().EdgeVisibilityOn()
  actor.SetPosition(kwargs['position'])

  return actor

def hex_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 5,
    'u_cyclic': True,
    'v_cyclic': False,
    'position': [0, 0, 0]
  }
  return tessellate(cylinder, HexTessagon, **options)

def tri_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 35,
    'v_num': 12,
    'position': [0, 15, 0]
  }
  return tessellate(torus, TriTessagon, **options)

def rhombus_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 6,
    'v_twist': True,
    'position': [0, 30, 0]
  }
  return tessellate(klein, RhombusTessagon, **options)

def octo_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 4,
    'v_num': 40,
    'v_cyclic': True,
    'u_cyclic': False,
    'u_twist': True,
    'position': [0, 45, 0]
  }
  return tessellate(mobius, OctoTessagon, **options)

def hex_tri_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 6,
    'position': [0, 60, 0]
  }
  return tessellate(torus, HexTriTessagon, **options)

def hex_square_tri_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 45,
    'v_num': 5,
    'position': [15, 0, 0]
  }
  return tessellate(torus, HexSquareTriTessagon, **options)

def square_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 15,
    'v_num': 4,
    'rot_factor': 2,
    'position': [15, 15, 0]
  }
  return tessellate(torus, SquareTessagon, **options)

def pythagorean_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 25,
    'v_num': 6,
    'position': [15, 30, 0]
  }
  return tessellate(torus, PythagoreanTessagon, **options)

def brick_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 15,
    'v_num': 3,
    'rot_factor': 3,
    'position': [15, 45, 0]
  }
  return tessellate(torus, BrickTessagon, **options)

def dodeca_tessagon():
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 25,
    'v_num': 3,
    'position': [15, 60, 0]
  }
  return tessellate(torus, DodecaTessagon, **options)

main()
<commit_msg>Update the VTK example.<commit_after>import os, sys
import vtk
from vtk.util.colors import tomato

this_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(this_dir + '/../..')

from tessagon.adaptors.vtk_adaptor import VtkAdaptor

from tessagon.types import *
from tessagon.misc.shapes import *

def main():
  ren = vtk.vtkRenderer()
  renWin = vtk.vtkRenderWindow()
  renWin.AddRenderer(ren)
  iren = vtk.vtkRenderWindowInteractor()
  iren.SetRenderWindow(renWin)

  ren.AddActor(hex_tessagon([0, 0, 0]))
  ren.AddActor(tri_tessagon([0, 15, 0]))
  ren.AddActor(rhombus_tessagon([0, 30, 0]))
  ren.AddActor(octo_tessagon([0, 45, 0]))
  ren.AddActor(hex_tri_tessagon([15, 0, 0]))
  ren.AddActor(hex_square_tri_tessagon([15, 15, 0]))
  ren.AddActor(square_tessagon([15, 30, 0]))
  ren.AddActor(pythagorean_tessagon([15, 45, 0]))
  ren.AddActor(brick_tessagon([30, 0, 0]))
  ren.AddActor(dodeca_tessagon([30, 15, 0]))
  ren.AddActor(square_tri_tessagon([30, 30, 0]))
  ren.AddActor(weave_tessagon([30, 45, 0]))

  ren.SetBackground(0.3, 0.3, 0.3)
  renWin.SetSize(800, 600)
  iren.Initialize()
  ren.ResetCamera()
  renWin.Render()
  iren.Start()

def tessellate(f, tessagon_class, **kwargs):
  extra_args = {'adaptor_class' : VtkAdaptor}
  tessagon = tessagon_class(f, **{**kwargs, **extra_args})

  poly_data = tessagon.create_mesh()
  mapper = vtk.vtkPolyDataMapper()
  mapper.SetInputData(poly_data)
  actor = vtk.vtkActor()
  actor.SetMapper(mapper)
  actor.GetProperty().SetColor(tomato)
  actor.GetProperty().SetEdgeColor(0.0, 0.0, 0.0)
  actor.GetProperty().EdgeVisibilityOn()
  actor.SetPosition(kwargs['position'])

  return actor

def hex_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 5,
    'u_cyclic': True,
    'v_cyclic': False,
    'position': position
  }
  return tessellate(cylinder, HexTessagon, **options)

def tri_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 35,
    'v_num': 12,
    'position': position
  }
  return tessellate(torus, TriTessagon, **options)

def rhombus_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 6,
    'v_twist': True,
    'position': position
  }
  return tessellate(klein, RhombusTessagon, **options)

def octo_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 4,
    'v_num': 40,
    'v_cyclic': True,
    'u_cyclic': False,
    'u_twist': True,
    'position': position
  }
  return tessellate(mobius, OctoTessagon, **options)

def hex_tri_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 40,
    'v_num': 6,
    'position': position
  }
  return tessellate(torus, HexTriTessagon, **options)

def hex_square_tri_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 45,
    'v_num': 5,
    'position': position
  }
  return tessellate(torus, HexSquareTriTessagon, **options)

def square_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 15,
    'v_num': 4,
    'rot_factor': 2,
    'position': position
  }
  return tessellate(torus, SquareTessagon, **options)

def pythagorean_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 25,
    'v_num': 6,
    'position': position
  }
  return tessellate(torus, PythagoreanTessagon, **options)

def brick_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 15,
    'v_num': 3,
    'rot_factor': 3,
    'position': position
  }
  return tessellate(torus, BrickTessagon, **options)

def dodeca_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 25,
    'v_num': 3,
    'position': position
  }
  return tessellate(torus, DodecaTessagon, **options)

def square_tri_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 20,
    'v_num': 4,
    'position': position
  }
  return tessellate(torus, SquareTriTessagon, **options)

def weave_tessagon(position):
  options = {
    'u_range': [0.0, 1.0],
    'v_range': [0.0, 1.0],
    'u_num': 8,
    'v_num': 6,
    'v_cyclic': False,
    'rot_factor': 1,
    'position': position
  }
  return tessellate(sphere, WeaveTessagon, **options)

main()
<|endoftext|>


<commit_before>import datetime
from decimal import Decimal
from operator import itemgetter
from optparse import OptionParser
import csv
import json
import math
from optparse import make_option
import sys
import urllib
import urllib2

from dateutil.parser import parse as dateparse

from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
from ngrams.models import *

class Calculator(object):
    def __init__(self, facet_field, congress=None):
        #self.n = n
        #self.gram = ['unigrams', 'bigrams', 'trigrams', 'quadgrams', 'pentagrams', ][self.n-1]
        self.facet_field = facet_field
        self.url = 'http://ec2-184-72-184-231.compute-1.amazonaws.com:8983/solr/select'
        self.df = {}
        self.congress = congress

    def gram(self, n):
        return ['unigrams', 'bigrams', 'trigrams', 'quadgrams', 'pentagrams', ][n-1]

    def _query_solr(self, data):
        body = urllib.urlencode(data)
        url = '%s?%s' % (self.url, body)
        r = urllib2.urlopen(url)
        return json.loads(r.read())

    def docs(self):
        if not hasattr(self, 'numdocs'):
            if self.congress:
                q = 'congress:"%s"' % self.congress
            else:
                q = '*:*'
            kwargs = {'q': q,
                      'facet': 'true',
                      'rows': 0,
                      'facet.field': self.facet_field,
                      'facet.limit': -1,
                      'facet.mincount': 0,
                      'facet.method': 'enum',
                      'wt': 'json', }
            data = self._query_solr(kwargs)
            self.numdocs = Decimal(len(data['facet_counts']['facet_fields'][self.facet_field]) / 2)

        return self.numdocs

    def docfreq(self, n, ngram):
        q = ['%s:"%s"' % (self.gram(n), ngram), ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        if not self.df.get(ngram):
            #kwargs = {'q': '%s:"%s"' % (self.gram(n), ngram),
            kwargs = {'q': ' AND '.join(q),
                      'facet': 'true',
                      'rows': 0,
                      'facet.field': self.facet_field,
                      'facet.limit': -1,
                      'facet.mincount': 1,
                      'facet.method': 'enum',
                      'wt': 'json', }
            data = self._query_solr(kwargs)
            self.df[ngram] = Decimal(len(data['facet_counts']['facet_fields'][self.facet_field]) / 2)
        return self.df.get(ngram)

    def list_facets(self):
        if self.congress:
            q = 'congress:"%s"' % self.congress
        else:
            q = '*:*'
        kwargs = {'q': q,
                  'facet': 'true',
                  'facet.field': self.facet_field,
                  'rows': 0,
                  'facet.method': 'enum',
                  'wt': 'json',
                  'facet.limit': -1,
                  'facet.mincount': 0, }
        data = self._query_solr(kwargs)
        return data['facet_counts']['facet_fields'][self.facet_field][::2]

    def list_ngrams_for_facet(self, n, facet):
        q = ['%s:"%s"' % (self.facet_field, facet), ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        #kwargs = {'q': '(%s:"%s") AND speaker_bioguide:[\'\' TO *]' % (self.facet_field, facet),
        kwargs = {'q': ' AND '.join(q),
                  'facet': 'true',
                  'facet.field': self.gram(n),
                  'facet.method': 'enumtermfreq',
                  'facet.mincount': 3,
                  'wt': 'json',
                  'rows': 0,
                  'facet.limit': 1000, }
        data = self._query_solr(kwargs)
        data = data['facet_counts']['facet_fields'][self.gram(n)]
        return zip(data[::2], data[1::2])

    def get_date_diversity(self, bioguide, n, ngram):
        url = 'http://localhost:8983/solr/select'
        q = ['speaker_bioguide:"%s"' % bioguide,
             '%s:"%s"' % (self.gram(n), ngram),
             ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        #kwargs = {'q': '(speaker_bioguide:"%s" AND %s:"%s")' % (bioguide, self.gram(n), ngram),
        kwargs = {'q': ' AND '.join(q),
                  'facet': 'true',
                  'facet.field': 'date',
                  'facet.method': 'enum',
                  'facet.mincount': 1,
                  'facet.limit': -1,
                  'wt': 'json', }
        data = self._query_solr(kwargs)
        data = data['facet_counts']['facet_fields']['date']
        return len(data)/2

def list_active_legislators_first():
    query = 'select bioguide_id from bioguide_legislatorrole group by bioguide_id order by max(end_date) desc, bioguide_id'
    from MySQLdb import Connection
    cursor = Connection('localhost', 'capwords', 'capwords', 'capwords').cursor()
    cursor.execute(query)
    results = [x[0] for x in cursor.fetchall()]
    return results

#if __name__ == '__main__':

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
            make_option('--field',
                action='store',
                dest='field',
                default=None,
                help='Field to calculate ngram tfidf for'),
            make_option('--congress',
                action='store',
                dest='congress',
                help='Restrict calculation to the given congress'),
            make_option('--values',
                action='store',
                dest='field_values',
                default=None,
                help='Specific values to iterate for field, if applicable'),
    )

    def handle(self, *args, **options):
        field = options.get('field')
        congress = options.get('congress')
        field_values = options.get('field_values')

        calculator = Calculator(field, congress)
        # generate facets based on somewhat sane defaults, or accept a comma-delimited list of strings to compare
        if field == 'speaker_bioguide':
            # already = set([(x[0], int(x[1])) for x in csv.reader(open(r'ngrams_by_bioguide.csv', 'r')) if len(x) > 1])
            already = set([(x.bioguide_id, int(x.n)) for x in NgramsByBioguide.objects.raw('select * from ngrams_ngramsbybioguide group by bioguide_id, n')])
            if field_values:
                facets = [facet.strip() for facet in field_values.split(',')]
            else:
                facets = list_active_legislators_first()
        elif field == 'speaker_state':
            already = set([(x.state, int(x.n)) for x in NgramsByState.objects.raw('select * from ngrams_ngramsbystate group by state, n')])
            if field_values:
                facets = [facet.strip() for facet in field_values.split(',')]
            else:
                facets = calculator.list_facets()
        elif field == 'year_month':
            # already = set([(x[0], int(x[1])) for x in csv.reader(open(r'ngrams_by_month.csv', 'r')) if len(x) > 1])
            already = set([(x.month, int(x.n)) for x in NgramsByMonth.objects.raw('select * from ngrams_ngramsbymonth group by month, n')])
            if field_values:
                facets = [int(facet.strip()) for facet in field_values.split(',')]
            else:
                facets = calculator.list_facets()
        elif field == 'date':
            already = []
            if field_values:
                missing = [dateparse(facet.strip()) for facet in field_values.split(',')]
            else:
                dates = Date.objects.values_list('date', flat=True).order_by('-date').distinct()[:1]
                missing = []
                for date in dates:
                    if NgramsByDate.objects.filter(date=date).count() == 0:
                        missing.append(date)
            facets = ['%sT12:00:00Z' % date.strftime('%Y-%m-%d') for date in missing]
        else:
            already = []
            facets = calculator.list_facets()

        for facet in reversed(facets):
            with transaction.commit_on_success():

                print facet
                for n in range(1,6):

                    if (facet, n) in already:
                        # recreate 'this' month, as of yesterday
                        if field == 'year_month' and facet == (datetime.datetime.today() - datetime.timedelta(1)).strftime('%Y%m'):
                            pass
                        # always recreate terms for speakers and states
                        elif field in ['speaker_bioguide', 'speaker_state']:
                            pass
                        elif field == 'date' and facet[0:10] in field_values:
                            pass
                        else:
                            continue

                    ngrams = []
                    ngrams_for_facet = calculator.list_ngrams_for_facet(n, facet)
                    facet_total_ngrams = sum([x[1] for x in ngrams_for_facet])

                    for ngram, count in ngrams_for_facet:
                        try:
                            df = calculator.docfreq(n, ngram)
                            idf = math.log(calculator.docs() / df)
                            tf = count / float(facet_total_ngrams)
                            tfidf = tf * idf
                            if tfidf == 0:
                                continue
                            #diversity = calculator.get_date_diversity(facet, n, ngram)
                            ngrams.append((ngram, tfidf, count))
                        except:
                            #print 'ERROR'
                            continue

                    ngrams.sort(key=itemgetter(1), reverse=True)

                    for ngram, tfidf, count in ngrams:
                        output = map(str, [facet, n, ngram, tfidf, int(count), ])
                        if congress:
                            output.append(congress)
                        #writer.writerow(output)
                        if field == 'date':
                            NgramsByDate.objects.filter(date=date, n=n, ngram=ngram).delete()
                            NgramsByDate.objects.create(n=n,
                                                        date=date,
                                                        ngram=ngram,
                                                        tfidf=tfidf,
                                                        count=count)
                        if field == 'year_month':
                            NgramsByMonth.objects.filter(month=facet, n=n, ngram=ngram).delete()
                            NgramsByMonth.objects.create(n=n,
                                                         month=facet,
                                                         ngram=ngram,
                                                         tfidf=tfidf,
                                                         count=count)
                        if field == 'speaker_bioguide':
                            NgramsByMonth.objects.filter(bioguide_id=facet, n=n, ngram=ngram).delete()
                            NgramsByBioguide.objects.create(n=n,
                                                            bioguide_id=facet,
                                                            ngram=ngram,
                                                            tfidf=tfidf,
                                                            count=count)
                        if field == 'speaker_state':
                            NgramsByState.objects.filter(state=facet, n=n, ngram=ngram).delete()
                            NgramsByState.objects.create(n=n,
                                                         state=facet,
                                                         ngram=ngram,
                                                         tfidf=tfidf,
                                                         count=count)
<commit_msg>whoo boy.<commit_after>import datetime
from decimal import Decimal
from operator import itemgetter
from optparse import OptionParser
import csv
import json
import math
from optparse import make_option
import sys
import urllib
import urllib2

from dateutil.parser import parse as dateparse

from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
from ngrams.models import *

class Calculator(object):
    def __init__(self, facet_field, congress=None):
        #self.n = n
        #self.gram = ['unigrams', 'bigrams', 'trigrams', 'quadgrams', 'pentagrams', ][self.n-1]
        self.facet_field = facet_field
        self.url = 'http://ec2-184-72-184-231.compute-1.amazonaws.com:8983/solr/select'
        self.df = {}
        self.congress = congress

    def gram(self, n):
        return ['unigrams', 'bigrams', 'trigrams', 'quadgrams', 'pentagrams', ][n-1]

    def _query_solr(self, data):
        body = urllib.urlencode(data)
        url = '%s?%s' % (self.url, body)
        r = urllib2.urlopen(url)
        return json.loads(r.read())

    def docs(self):
        if not hasattr(self, 'numdocs'):
            if self.congress:
                q = 'congress:"%s"' % self.congress
            else:
                q = '*:*'
            kwargs = {'q': q,
                      'facet': 'true',
                      'rows': 0,
                      'facet.field': self.facet_field,
                      'facet.limit': -1,
                      'facet.mincount': 0,
                      'facet.method': 'enum',
                      'wt': 'json', }
            data = self._query_solr(kwargs)
            self.numdocs = Decimal(len(data['facet_counts']['facet_fields'][self.facet_field]) / 2)

        return self.numdocs

    def docfreq(self, n, ngram):
        q = ['%s:"%s"' % (self.gram(n), ngram), ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        if not self.df.get(ngram):
            #kwargs = {'q': '%s:"%s"' % (self.gram(n), ngram),
            kwargs = {'q': ' AND '.join(q),
                      'facet': 'true',
                      'rows': 0,
                      'facet.field': self.facet_field,
                      'facet.limit': -1,
                      'facet.mincount': 1,
                      'facet.method': 'enum',
                      'wt': 'json', }
            data = self._query_solr(kwargs)
            self.df[ngram] = Decimal(len(data['facet_counts']['facet_fields'][self.facet_field]) / 2)
        return self.df.get(ngram)

    def list_facets(self):
        if self.congress:
            q = 'congress:"%s"' % self.congress
        else:
            q = '*:*'
        kwargs = {'q': q,
                  'facet': 'true',
                  'facet.field': self.facet_field,
                  'rows': 0,
                  'facet.method': 'enum',
                  'wt': 'json',
                  'facet.limit': -1,
                  'facet.mincount': 0, }
        data = self._query_solr(kwargs)
        return data['facet_counts']['facet_fields'][self.facet_field][::2]

    def list_ngrams_for_facet(self, n, facet):
        q = ['%s:"%s"' % (self.facet_field, facet), ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        #kwargs = {'q': '(%s:"%s") AND speaker_bioguide:[\'\' TO *]' % (self.facet_field, facet),
        kwargs = {'q': ' AND '.join(q),
                  'facet': 'true',
                  'facet.field': self.gram(n),
                  'facet.method': 'enumtermfreq',
                  'facet.mincount': 3,
                  'wt': 'json',
                  'rows': 0,
                  'facet.limit': 1000, }
        data = self._query_solr(kwargs)
        data = data['facet_counts']['facet_fields'][self.gram(n)]
        return zip(data[::2], data[1::2])

    def get_date_diversity(self, bioguide, n, ngram):
        url = 'http://localhost:8983/solr/select'
        q = ['speaker_bioguide:"%s"' % bioguide,
             '%s:"%s"' % (self.gram(n), ngram),
             ]
        if self.congress:
            q.append('congress:"%s"' % self.congress)

        #kwargs = {'q': '(speaker_bioguide:"%s" AND %s:"%s")' % (bioguide, self.gram(n), ngram),
        kwargs = {'q': ' AND '.join(q),
                  'facet': 'true',
                  'facet.field': 'date',
                  'facet.method': 'enum',
                  'facet.mincount': 1,
                  'facet.limit': -1,
                  'wt': 'json', }
        data = self._query_solr(kwargs)
        data = data['facet_counts']['facet_fields']['date']
        return len(data)/2

def list_active_legislators_first():
    query = 'select bioguide_id from bioguide_legislatorrole group by bioguide_id order by max(end_date) desc, bioguide_id'
    from MySQLdb import Connection
    cursor = Connection('localhost', 'capwords', 'capwords', 'capwords').cursor()
    cursor.execute(query)
    results = [x[0] for x in cursor.fetchall()]
    return results

#if __name__ == '__main__':

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
            make_option('--field',
                action='store',
                dest='field',
                default=None,
                help='Field to calculate ngram tfidf for'),
            make_option('--congress',
                action='store',
                dest='congress',
                help='Restrict calculation to the given congress'),
            make_option('--values',
                action='store',
                dest='field_values',
                default=None,
                help='Specific values to iterate for field, if applicable'),
    )

    def handle(self, *args, **options):
        field = options.get('field')
        congress = options.get('congress')
        field_values = options.get('field_values')

        calculator = Calculator(field, congress)
        # generate facets based on somewhat sane defaults, or accept a comma-delimited list of strings to compare
        if field == 'speaker_bioguide':
            # already = set([(x[0], int(x[1])) for x in csv.reader(open(r'ngrams_by_bioguide.csv', 'r')) if len(x) > 1])
            already = set([(x.bioguide_id, int(x.n)) for x in NgramsByBioguide.objects.raw('select * from ngrams_ngramsbybioguide group by bioguide_id, n')])
            if field_values:
                facets = [facet.strip() for facet in field_values.split(',')]
            else:
                facets = list_active_legislators_first()
        elif field == 'speaker_state':
            already = set([(x.state, int(x.n)) for x in NgramsByState.objects.raw('select * from ngrams_ngramsbystate group by state, n')])
            if field_values:
                facets = [facet.strip() for facet in field_values.split(',')]
            else:
                facets = calculator.list_facets()
        elif field == 'year_month':
            # already = set([(x[0], int(x[1])) for x in csv.reader(open(r'ngrams_by_month.csv', 'r')) if len(x) > 1])
            already = set([(x.month, int(x.n)) for x in NgramsByMonth.objects.raw('select * from ngrams_ngramsbymonth group by month, n')])
            if field_values:
                facets = [int(facet.strip()) for facet in field_values.split(',')]
            else:
                facets = calculator.list_facets()
        elif field == 'date':
            already = []
            if field_values:
                missing = [dateparse(facet.strip()) for facet in field_values.split(',')]
            else:
                dates = Date.objects.values_list('date', flat=True).order_by('-date').distinct()[:1]
                missing = []
                for date in dates:
                    if NgramsByDate.objects.filter(date=date).count() == 0:
                        missing.append(date)
            facets = ['%sT12:00:00Z' % date.strftime('%Y-%m-%d') for date in missing]
        else:
            already = []
            facets = calculator.list_facets()

        for facet in reversed(facets):
            with transaction.commit_on_success():

                print facet
                for n in range(1,6):

                    if (facet, n) in already:
                        # recreate 'this' month, as of yesterday
                        if field == 'year_month' and facet == (datetime.datetime.today() - datetime.timedelta(1)).strftime('%Y%m'):
                            pass
                        # always recreate terms for speakers and states
                        elif field in ['speaker_bioguide', 'speaker_state']:
                            pass
                        elif field == 'date' and facet[0:10] in field_values:
                            pass
                        else:
                            continue

                    ngrams = []
                    ngrams_for_facet = calculator.list_ngrams_for_facet(n, facet)
                    facet_total_ngrams = sum([x[1] for x in ngrams_for_facet])

                    for ngram, count in ngrams_for_facet:
                        try:
                            df = calculator.docfreq(n, ngram)
                            idf = math.log(calculator.docs() / df)
                            tf = count / float(facet_total_ngrams)
                            tfidf = tf * idf
                            if tfidf == 0:
                                continue
                            #diversity = calculator.get_date_diversity(facet, n, ngram)
                            ngrams.append((ngram, tfidf, count))
                        except:
                            #print 'ERROR'
                            continue

                    ngrams.sort(key=itemgetter(1), reverse=True)

                    for ngram, tfidf, count in ngrams:
                        output = map(str, [facet, n, ngram, tfidf, int(count), ])
                        if congress:
                            output.append(congress)
                        #writer.writerow(output)
                        if field == 'date':
                            NgramsByDate.objects.filter(date=date, n=n, ngram=ngram).delete()
                            NgramsByDate.objects.create(n=n,
                                                        date=date,
                                                        ngram=ngram,
                                                        tfidf=tfidf,
                                                        count=count)
                        if field == 'year_month':
                            NgramsByMonth.objects.filter(month=facet, n=n, ngram=ngram).delete()
                            NgramsByMonth.objects.create(n=n,
                                                         month=facet,
                                                         ngram=ngram,
                                                         tfidf=tfidf,
                                                         count=count)
                        if field == 'speaker_bioguide':
                            NgramsByBioguide.objects.filter(bioguide_id=facet, n=n, ngram=ngram).delete()
                            NgramsByBioguide.objects.create(n=n,
                                                            bioguide_id=facet,
                                                            ngram=ngram,
                                                            tfidf=tfidf,
                                                            count=count)
                        if field == 'speaker_state':
                            NgramsByState.objects.filter(state=facet, n=n, ngram=ngram).delete()
                            NgramsByState.objects.create(n=n,
                                                         state=facet,
                                                         ngram=ngram,
                                                         tfidf=tfidf,
                                                         count=count)
<|endoftext|>