<commit_before>__author__ = 'Bohdan Mushkevych'

import collections
from threading import Lock
from datetime import datetime, timedelta

from synergy.conf import settings
from synergy.system.data_logging import get_logger
from synergy.system.decorator import thread_safe
from synergy.system.priority_queue import PriorityEntry, PriorityQueue, compute_release_time
from synergy.system.repeat_timer import RepeatTimer
from synergy.scheduler.scheduler_constants import QUEUE_UOW_REPORT, PROCESS_GC
from synergy.scheduler.thread_handler import ManagedThreadHandler
from synergy.db.model import unit_of_work
from synergy.db.model.synergy_mq_transmission import SynergyMqTransmission
from synergy.db.dao.unit_of_work_dao import UnitOfWorkDao

class GarbageCollector(object):
    """ GC is triggered directly by Synergy Scheduler.
        It scans for invalid or stalled unit_of_work and re-triggers them.
        GC is vital for the health of the system.
        Deployment with no GC is considered invalid """

    def __init__(self, scheduler):
        self.logger = get_logger(PROCESS_GC, append_to_console=False, redirect_stdstream=False)
        self.managed_handlers = scheduler.managed_handlers
        self.publishers = scheduler.publishers
        self.timetable = scheduler.timetable

        self.lock = Lock()
        self.uow_dao = UnitOfWorkDao(self.logger)
        self.reprocess_uows = collections.defaultdict(PriorityQueue)
        self.timer = RepeatTimer(settings.settings['gc_run_interval'], self._run)

    @thread_safe
    def scan_uow_candidates(self):
        """ method performs two actions:
            - enlist stale or invalid units of work into reprocessing queue
            - cancel UOWs that are older than 2 days and have been submitted more than 1 hour ago """
        try:
            since = settings.settings['synergy_start_timeperiod']
            uow_list = self.uow_dao.get_reprocessing_candidates(since)
        except LookupError as e:
            self.logger.info('flow: re-processing UOW candidates not found. %r' % e)
            return

        for uow in uow_list:
            try:
                if uow.process_name not in self.managed_handlers:
                    self.logger.debug('process %r is not known to the Synergy Scheduler. Skipping its unit_of_work.'
                                      % uow.process_name)
                    continue

                thread_handler = self.managed_handlers[uow.process_name]
                assert isinstance(thread_handler, ManagedThreadHandler)

                if not thread_handler.process_entry.is_on:
                    self.logger.debug('process %r is inactive. Skipping its unit_of_work.' % uow.process_name)
                    continue

                entry = PriorityEntry(uow)
                if entry in self.reprocess_uows[uow.process_name]:
                    # given UOW is already registered in the reprocessing queue
                    continue

                # ASSUMPTION: UOW is re-created by a state machine during reprocessing
                # thus - any UOW older 2 days could be marked as STATE_CANCELED
                if datetime.utcnow() - uow.created_at > timedelta(hours=settings.settings['gc_life_support_hours']):
                    self._cancel_uow(uow)
                    continue

                # if the UOW has been idle for more than 1 hour - resubmit it
                if datetime.utcnow() - uow.submitted_at > timedelta(hours=settings.settings['gc_resubmit_after_hours'])\
                        or uow.is_invalid:
                    # enlist the UOW into the reprocessing queue
                    self.reprocess_uows[uow.process_name].put(entry)

            except Exception as e:
                self.logger.error('flow exception: %s' % str(e), exc_info=True)

    def _flush_queue(self, q, ignore_priority=False):
        """
        :param q: PriorityQueue instance holding GarbageCollector entries
        :param ignore_priority: If True - all GarbageCollector entries should be resubmitted
                If False - only those entries whose waiting time has expired will be resubmitted
        """
        assert isinstance(q, PriorityQueue)

        current_timestamp = compute_release_time(lag_in_minutes=0)
        for _ in range(len(q)):
            entry = q.pop()
            assert isinstance(entry, PriorityEntry)

            if ignore_priority or entry.release_time < current_timestamp:
                self._resubmit_uow(entry.entry)
            else:
                q.put(entry)
                break

    @thread_safe
    def flush(self, ignore_priority=False):
        """ method iterates over each reprocessing queues and re-submits UOW whose waiting time has expired """
        for process_name, q in self.reprocess_uows.items():
            self._flush_queue(q, ignore_priority)

    @thread_safe
    def flush_one(self, process_name, ignore_priority=False):
        """ method iterates over the reprocessing queue for the given process
            and re-submits UOW whose waiting time has expired """
        q = self.reprocess_uows[process_name]
        self._flush_queue(q, ignore_priority)

    def _resubmit_uow(self, uow):
        mq_request = SynergyMqTransmission(process_name=uow.process_name, unit_of_work_id=uow.db_id)

        if uow.is_invalid:
            uow.number_of_retries += 1

        uow.state = unit_of_work.STATE_REQUESTED
        uow.submitted_at = datetime.utcnow()
        self.uow_dao.update(uow)

        publisher = self.publishers.get(uow.process_name)
        publisher.publish(mq_request.document)
        publisher.release()

        self.logger.info('re-submitted UOW {0} for {1} in {2}; attempt {3}'
                         .format(uow.db_id, uow.process_name, uow.timeperiod, uow.number_of_retries))

    def _cancel_uow(self, uow):
        mq_request = SynergyMqTransmission(process_name=uow.process_name, unit_of_work_id=uow.db_id)

        uow.state = unit_of_work.STATE_CANCELED
        self.uow_dao.update(uow)

        publisher = self.publishers.get(QUEUE_UOW_REPORT)
        publisher.publish(mq_request.document)
        publisher.release()

        self.logger.info('canceled UOW {0} for {1} in {2}; attempt {3}; created at {4}'
                         .format(uow.db_id, uow.process_name, uow.timeperiod, uow.number_of_retries, uow.created_at))

    def _run(self):
        try:
            self.logger.info('run {')

            self.logger.debug('GC: step 1 - enlist or cancel')
            self.scan_uow_candidates()

            self.logger.debug('GC: step 2 - repost after timeout')
            self.flush()

            self.logger.debug('GC: step 3 - timetable housekeeping')
            self.timetable.build_trees()

            self.logger.debug('GC: step 4 - timetable validation')
            self.timetable.validate()
        except Exception as e:
            self.logger.error('GC run exception: %s' % str(e))
        finally:
            self.logger.info('}')

    def start(self):
        self.timer.start()

    def stop(self):
        self.timer.cancel()
<commit_msg>- logger messages rewording<commit_after>__author__ = 'Bohdan Mushkevych'

import collections
from threading import Lock
from datetime import datetime, timedelta

from synergy.conf import settings
from synergy.system.data_logging import get_logger
from synergy.system.decorator import thread_safe
from synergy.system.priority_queue import PriorityEntry, PriorityQueue, compute_release_time
from synergy.system.repeat_timer import RepeatTimer
from synergy.scheduler.scheduler_constants import QUEUE_UOW_REPORT, PROCESS_GC
from synergy.scheduler.thread_handler import ManagedThreadHandler
from synergy.db.model import unit_of_work
from synergy.db.model.synergy_mq_transmission import SynergyMqTransmission
from synergy.db.dao.unit_of_work_dao import UnitOfWorkDao

class GarbageCollector(object):
    """ GC is triggered directly by Synergy Scheduler.
        It scans for invalid or stalled unit_of_work and re-triggers them.
        GC is vital for the health of the system.
        Deployment with no GC is considered invalid """

    def __init__(self, scheduler):
        self.logger = get_logger(PROCESS_GC, append_to_console=False, redirect_stdstream=False)
        self.managed_handlers = scheduler.managed_handlers
        self.publishers = scheduler.publishers
        self.timetable = scheduler.timetable

        self.lock = Lock()
        self.uow_dao = UnitOfWorkDao(self.logger)
        self.reprocess_uows = collections.defaultdict(PriorityQueue)
        self.timer = RepeatTimer(settings.settings['gc_run_interval'], self._run)

    @thread_safe
    def scan_uow_candidates(self):
        """ method performs two actions:
            - enlist stale or invalid units of work into reprocessing queue
            - cancel UOWs that are older than 2 days and have been submitted more than 1 hour ago """
        try:
            since = settings.settings['synergy_start_timeperiod']
            uow_list = self.uow_dao.get_reprocessing_candidates(since)
        except LookupError as e:
            self.logger.info('flow: no UOW candidates found for reprocessing: {0}'.format(e))
            return

        for uow in uow_list:
            try:
                if uow.process_name not in self.managed_handlers:
                    self.logger.debug('process {0} is not known to the Synergy Scheduler. Skipping its UOW.'
                                      .format(uow.process_name))
                    continue

                thread_handler = self.managed_handlers[uow.process_name]
                assert isinstance(thread_handler, ManagedThreadHandler)

                if not thread_handler.process_entry.is_on:
                    self.logger.debug('process {0} is inactive. Skipping its UOW.'.format(uow.process_name))
                    continue

                entry = PriorityEntry(uow)
                if entry in self.reprocess_uows[uow.process_name]:
                    # given UOW is already registered in the reprocessing queue
                    continue

                # ASSUMPTION: UOW is re-created by a state machine during reprocessing
                # thus - any UOW older 2 days could be marked as STATE_CANCELED
                if datetime.utcnow() - uow.created_at > timedelta(hours=settings.settings['gc_life_support_hours']):
                    self._cancel_uow(uow)
                    continue

                # if the UOW has been idle for more than 1 hour - resubmit it
                if datetime.utcnow() - uow.submitted_at > timedelta(hours=settings.settings['gc_resubmit_after_hours'])\
                        or uow.is_invalid:
                    # enlist the UOW into the reprocessing queue
                    self.reprocess_uows[uow.process_name].put(entry)

            except Exception as e:
                self.logger.error('flow exception: {0}'.format(e), exc_info=True)

    def _flush_queue(self, q, ignore_priority=False):
        """
        :param q: PriorityQueue instance holding GarbageCollector entries
        :param ignore_priority: If True - all GarbageCollector entries should be resubmitted
                If False - only those entries whose waiting time has expired will be resubmitted
        """
        assert isinstance(q, PriorityQueue)

        current_timestamp = compute_release_time(lag_in_minutes=0)
        for _ in range(len(q)):
            entry = q.pop()
            assert isinstance(entry, PriorityEntry)

            if ignore_priority or entry.release_time < current_timestamp:
                self._resubmit_uow(entry.entry)
            else:
                q.put(entry)
                break

    @thread_safe
    def flush(self, ignore_priority=False):
        """ method iterates over each reprocessing queues and re-submits UOW whose waiting time has expired """
        for process_name, q in self.reprocess_uows.items():
            self._flush_queue(q, ignore_priority)

    @thread_safe
    def flush_one(self, process_name, ignore_priority=False):
        """ method iterates over the reprocessing queue for the given process
            and re-submits UOW whose waiting time has expired """
        q = self.reprocess_uows[process_name]
        self._flush_queue(q, ignore_priority)

    def _resubmit_uow(self, uow):
        mq_request = SynergyMqTransmission(process_name=uow.process_name, unit_of_work_id=uow.db_id)

        if uow.is_invalid:
            uow.number_of_retries += 1

        uow.state = unit_of_work.STATE_REQUESTED
        uow.submitted_at = datetime.utcnow()
        self.uow_dao.update(uow)

        publisher = self.publishers.get(uow.process_name)
        publisher.publish(mq_request.document)
        publisher.release()

        self.logger.info('re-submitted UOW {0} for {1} in {2}; attempt {3}'
                         .format(uow.db_id, uow.process_name, uow.timeperiod, uow.number_of_retries))

    def _cancel_uow(self, uow):
        mq_request = SynergyMqTransmission(process_name=uow.process_name, unit_of_work_id=uow.db_id)

        uow.state = unit_of_work.STATE_CANCELED
        self.uow_dao.update(uow)

        publisher = self.publishers.get(QUEUE_UOW_REPORT)
        publisher.publish(mq_request.document)
        publisher.release()

        self.logger.info('canceled UOW {0} for {1} in {2}; attempt {3}; created at {4}'
                         .format(uow.db_id, uow.process_name, uow.timeperiod, uow.number_of_retries, uow.created_at))

    def _run(self):
        try:
            self.logger.info('run {')

            self.logger.debug('GC: step 1 - scan reprocessing candidates')
            self.scan_uow_candidates()

            self.logger.debug('GC: step 2 - repost after timeout')
            self.flush()

            self.logger.debug('GC: step 3 - timetable housekeeping')
            self.timetable.build_trees()

            self.logger.debug('GC: step 4 - timetable validation')
            self.timetable.validate()
        except Exception as e:
            self.logger.error('GC run exception: {0}'.format(e))
        finally:
            self.logger.info('}')

    def start(self):
        self.timer.start()

    def stop(self):
        self.timer.cancel()
<|endoftext|>


<commit_before>import xmlrpclib
import traceback
from cStringIO import StringIO

allowed = ('package_releases', 'package_urls', 'package_data',
    'search', 'list_packages', 'release_urls', 'release_data',
    'updated_releases', 'changelog', 'post_cheesecake_for_release')

def handle_request(webui_obj):
    webui_obj.handler.send_response(200, 'OK')
    webui_obj.handler.send_header('Content-type', 'text/xml')
    webui_obj.handler.send_header('charset', 'UTF-8' );
    webui_obj.handler.end_headers()
    try:
        methodArgs, methodName = xmlrpclib.loads(webui_obj.handler.rfile.read())
        if methodName in allowed:
            response = globals()[methodName](webui_obj.store, *methodArgs)
        else:
            raise KeyError, "Method %r does not exist" % (methodName,)
        if response is None:
            response = ''
        # xmlrpclib.dumps encodes Unicode as UTF-8
        xml = xmlrpclib.dumps((response,), methodresponse=True, allow_none=True)
        webui_obj.handler.wfile.write(xml)
    except:
        out = StringIO()
        traceback.print_exc(file=out)
        result = xmlrpclib.dumps(xmlrpclib.Fault(1, out.getvalue()), methodresponse=True)
        webui_obj.handler.wfile.write(result)

def list_packages(store):
    result = store.get_packages()
    return [row['name'] for row in result]

def package_releases(store, package_name, show_hidden=False):
    if show_hidden:
        hidden = None
    else:
        hidden = False
    result = store.get_package_releases(package_name, hidden=hidden)
    return [row['version'] for row in result]

def release_urls(store, package_name, version):
    result = []
    for file in store.list_files(package_name, version):
        info = file.as_dict()
        info['url'] = store.gen_file_url(info['python_version'],
            package_name, info['filename'])
        result.append(info)
    # TODO do something with release_urls when there is something to do
    #info = store.get_package(package_name, version)
    #if info['download_url']:
    #    result.append({'url': info['download_url']})
    return result
package_urls = release_urls     # "deprecated"

def release_data(store, package_name, version):
    info = store.get_package(package_name, version).as_dict()
    del info['description_html']
    for col in ('requires', 'provides', 'obsoletes'):
        rows = store.get_release_relationships(package_name, version, col)
        info[col] = [row['specifier'] for row in rows]
    classifiers = [r[0] for r in store.get_release_classifiers(package_name,
        version)]
    info['classifiers' ] = classifiers
    return info
package_data = release_data     # "deprecated"

def search(store, spec, operator='and'):
    spec['_pypi_hidden'] = 'FALSE'
    return [row.as_dict() for row in store.query_packages(spec, operator)]

def updated_releases(store, since):
    result = store.updated_releases(since)
    return [(row['name'], row['version']) for row in result]

def changelog(store, since):
    result = store.changelog(since)
    return [(row['name'],row['version'],
             int(row['submitted_date'].gmticks()),
             row['action'])
             for row in result]

def post_cheesecake_for_release(store, name, version, score_data, password):
    if password != store.config.cheesecake_password:
        raise ValuError("Bad password.")

    store.save_cheesecake_score(name, version, score_data)
    store.commit()
<commit_msg>Support XML-RPC marshalling of mx.DateTime.<commit_after>import xmlrpclib
import traceback
from cStringIO import StringIO

allowed = ('package_releases', 'package_urls', 'package_data',
    'search', 'list_packages', 'release_urls', 'release_data',
    'updated_releases', 'changelog', 'post_cheesecake_for_release')

# monkey-patch xmlrpclib to marshal mx.DateTime correctly.
import mx.DateTime
def dump_DateTime(self, value, write):
    write("<value><dateTime.iso8601>")
    write(value.strftime("%Y%m%dT%H:%M:%S"))
    write("</dateTime.iso8601></value>\n")
xmlrpclib.Marshaller.dispatch[mx.DateTime.DateTimeType] = dump_DateTime

def handle_request(webui_obj):
    webui_obj.handler.send_response(200, 'OK')
    webui_obj.handler.send_header('Content-type', 'text/xml')
    webui_obj.handler.send_header('charset', 'UTF-8' );
    webui_obj.handler.end_headers()
    try:
        methodArgs, methodName = xmlrpclib.loads(webui_obj.handler.rfile.read())
        if methodName in allowed:
            response = globals()[methodName](webui_obj.store, *methodArgs)
        else:
            raise KeyError, "Method %r does not exist" % (methodName,)
        if response is None:
            response = ''
        # xmlrpclib.dumps encodes Unicode as UTF-8
        xml = xmlrpclib.dumps((response,), methodresponse=True, allow_none=True)
        webui_obj.handler.wfile.write(xml)
    except:
        out = StringIO()
        traceback.print_exc(file=out)
        result = xmlrpclib.dumps(xmlrpclib.Fault(1, out.getvalue()), methodresponse=True)
        webui_obj.handler.wfile.write(result)

def list_packages(store):
    result = store.get_packages()
    return [row['name'] for row in result]

def package_releases(store, package_name, show_hidden=False):
    if show_hidden:
        hidden = None
    else:
        hidden = False
    result = store.get_package_releases(package_name, hidden=hidden)
    return [row['version'] for row in result]

def release_urls(store, package_name, version):
    result = []
    for file in store.list_files(package_name, version):
        info = file.as_dict()
        info['url'] = store.gen_file_url(info['python_version'],
            package_name, info['filename'])
        result.append(info)
    # TODO do something with release_urls when there is something to do
    #info = store.get_package(package_name, version)
    #if info['download_url']:
    #    result.append({'url': info['download_url']})
    return result
package_urls = release_urls     # "deprecated"

def release_data(store, package_name, version):
    info = store.get_package(package_name, version).as_dict()
    del info['description_html']
    for col in ('requires', 'provides', 'obsoletes'):
        rows = store.get_release_relationships(package_name, version, col)
        info[col] = [row['specifier'] for row in rows]
    classifiers = [r[0] for r in store.get_release_classifiers(package_name,
        version)]
    info['classifiers' ] = classifiers
    return info
package_data = release_data     # "deprecated"

def search(store, spec, operator='and'):
    spec['_pypi_hidden'] = 'FALSE'
    return [row.as_dict() for row in store.query_packages(spec, operator)]

def updated_releases(store, since):
    result = store.updated_releases(since)
    return [(row['name'], row['version']) for row in result]

def changelog(store, since):
    result = store.changelog(since)
    return [(row['name'],row['version'],
             int(row['submitted_date'].gmticks()),
             row['action'])
             for row in result]

def post_cheesecake_for_release(store, name, version, score_data, password):
    if password != store.config.cheesecake_password:
        raise ValuError("Bad password.")

    store.save_cheesecake_score(name, version, score_data)
    store.commit()
<|endoftext|>


<commit_before>import asyncio
from concurrent import futures
import os
import random
import re
from urllib.parse import parse_qs

import discord
from discord.ext import commands
from lxml import etree

from starwars import starroller, die_names

class RPG:
    def __init__(self, bot):
        self.bot = bot
        self.tarot_url = 'http://www.tarotlore.com/tarot-cards/{}/'

    @commands.command()
    async def choose(self, ctx, *options):
        """Choose between some options. Use quotes if they have spaces."""
        await ctx.send(random.choice(options))

    @commands.command()
    async def tarot(self, ctx, *suits):
        """Get a random tarot card.

        You can specify the suits from which to pull, options are:
        minor:
            cups
            swords
            wands
            pentacles
        major"""
        async with ctx.typing():
            cards = []
            if not suits:
                suits = ('cups', 'swords', 'wands', 'pentacles', 'major')
            if 'minor' in suits:
                suits = suits + ('cups', 'swords', 'wands', 'pentacles')
            suits = set(suit.lower() for suit in suits)
            for root, dirs, files in os.walk('data/tarot'):
                if any(suit in root for suit in suits):
                    cards += [f'{root}/{card}' for card in files]
            try:
                card = random.choice(cards).replace('\\', '/')
            except IndexError:
                await ctx.send('Please specify a valid suit, or no suit.')
                return
            match = re.match(r'(?:\w+\/)+[IVX0_]*([\w_]+)\.jpg', card)
            name = match.groups()[0].replace('_', ' ')
            url = self.tarot_url.format(name.lower().replace(' ', '-'))
            embed = discord.Embed()
            embed.title = name
            embed.url = url
            filename = card.replace('_', '').rpartition('/')[-1]
            embed.set_image(url=f'attachment://{filename}')
            await ctx.send(file=discord.File(f'{card}', filename),
                           embed=embed)

    @commands.command(aliases=['r'])
    async def roll(self, ctx, *, inp='1d20'):
        """Roll some dice!

        Can roll multiple dice of any size, with modifiers.
        Format: XdY([+-^v]Z)(xN)(s)(t)
        X is the number of dice
        Y is the number of sides
        + adds Z to the result
        - subtracts Z from the result
        ^ drops the Z highest dice
        v drops the Z lowest dice
        x repeats the roll N times
        s sorts the results
        t totals each roll
        You can join rolls with commas
        """
        if inp == 'stats':
            inp = '4d6v1x6t'
        inp = ''.join(inp.split()).lower()
        expr = r'^([\d]*d?\d+([+\-^v]\d+)?(x\d+)?([ts]{1,2})?,?)+(?<!,)$'
        if re.match(expr, inp) is None:
            raise commands.BadArgument

        rolls = inp.split(',')
        args_batch = []
        for roll in rolls:
            if 'd' not in roll:
                roll = f'1d{roll}'
            elif roll[0] == 'd':
                roll = f'1{roll}'
            args = tuple(int(arg) for arg in re.findall(r'\d+', roll))

            num = args[0]
            sides = args[1]

            hi_drop = 0
            lo_drop = 0
            mod = 0

            if '^' in inp:
                hi_drop = args[2]
            elif 'v' in inp:
                lo_drop = args[2]
            elif '+' in inp:
                mod = args[2]
            elif '-' in inp:
                mod = -args[2]

            if 'x' in inp:
                times = args[-1]
            else:
                times = 1

            loop = self.bot.loop
            args = (num, sides, lo_drop, hi_drop, mod, times)
            args_batch.append(args)

        future = loop.run_in_executor(None, self.roll_helper, args_batch)
        async with ctx.typing():
            results = await asyncio.wait_for(future, 10, loop=loop)

        out = []
        for roll, result in zip(rolls, results):
            total = 't' in roll
            if total:
                result = [[sum(roll_)] for roll_ in result]
            if 's' in inp:
                for roll_ in result:
                    roll_.sort()
                result.sort()
            if total or len(result) == 1:
                result = [roll_[0] for roll_ in result]
            if times == 1:
                result = result[0]
            out.append(f'{roll}: {result}')
        await ctx.reply('\n'.join(out))

    def roll_helper(self, rolls):
        out = []
        for roll in rolls:
            out.append(roller(*roll))
        return out

    @roll.error
    async def roll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, (commands.MissingRequiredArgument,
                      commands.BadArgument)):
            await ctx.send('Invalid input. Valid input examples:'
                           '\n1d20+3'
                           '\n1d6'
                           '\n2d8-4'
                           '\n2d20^1'
                           '\n4d6v1x6t')
        elif isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        elif isinstance(e, discord.HTTPException):
            await ctx.reply('Your results were too long. Maybe sum them?')
        else:
            await self.bot.handle_error(ctx, e)

    @commands.command(aliases=['shadroll', 'sr'])
    async def shadowroll(self, ctx, *, inp):
        """Roll some dice - for Shadowrun!

        Format: N[e]
        Roll N six-sided dice and return the number of dice that rolled 5 or 6.
        If you put "e" after the number, 6s are counted and then rerolled."""
        inp = inp.strip()
        expr = r'^\d+e?$'
        if not re.match(expr, inp):
            raise commands.BadArgument
        edge = 'e' in inp
        if edge:
            inp = inp[:-1]
        num = int(inp)

        loop = self.bot.loop
        args = (num, edge)
        future = loop.run_in_executor(None, shadowroller, *args)
        async with ctx.typing():
            result = await asyncio.wait_for(future, 10, loop=loop)

        await ctx.reply(result)

    @shadowroll.error
    async def shadowroll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, (commands.MissingRequiredArgument,
                      commands.BadArgument)):
            await ctx.send('Invalid input. Valid input examples:'
                           '\n6'
                           '\n13e')
        elif isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        else:
            await self.bot.handle_error(ctx, e)

    @commands.command(aliases=['sw'])
    async def starroll(self, ctx, *, inp):
        """Roll some dice - for Fantasy Flight Star Wars!

        Available dice:
        b[oost]
        a[bility]
        p[roficiency]
        s[etback]
        d[ifficulty]
        c[hallenge]
        f[orce]

        Input examples:
        4a3d
        3a2p1b4d1c
        2f"""
        inp = inp.lower()
        expr = r'^(\d+[a-z])+$'
        match = re.match(expr, inp)
        if not match:
            raise commands.BadArgument
        dice = {}
        for roll in match.groups():
            num = int(roll[:-1])
            try:
                die = die_names[roll[-1]]
            except KeyError:
                await ctx.send(f'Die "{roll[-1]}" does not exist.')
                return
            dice[die] = num

        loop = self.bot.loop
        future = loop.run_in_executor(None, lambda: starroller(**dice))
        async with ctx.typing():
            try:
                result = await asyncio.wait_for(future, 10, loop=loop)
            except ValueError:
                await ctx.send('Force dice cannot be used with other dice.')
            else:
                await ctx.reply(result)

    @starroll.error
    async def starroll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        else:
            await self.bot.handle_error(ctx, e)

def roller(num=1, sides=20, lo_drop=0, hi_drop=0, mod=0, times=1):
    rolls = []
    for _ in range(times):
        pool = [random.randint(1, sides) for _ in range(num)]
        if lo_drop or hi_drop:
            sorted_pool = sorted(pool)
            dropped_vals = sorted_pool[:lo_drop] + sorted_pool[num-hi_drop:]
            for val in dropped_vals:
                pool.remove(val)
        if mod:
            pool = [sum(pool) + mod]
        rolls.append(pool)
    return rolls

def shadowroller(num, edge=False):
    rolls = hits = count1 = 0
    while True:
        count6 = 0
        rolls += num
        for _ in range(num):
            roll = random.randint(1, 6)
            if roll > 4:
                hits += 1
                if roll == 6:
                    count6 += 1
            elif roll == 1:
                count1 += 1
        if not (count6 > 0 and edge):
            break
        num = count6
    s = 's' if hits != 1 else ''
    if count1 < rolls / 2:
        result = f'{hits} hit{s}.'
    elif hits == 0:
        result = 'Critical glitch.'
    else:
        result = f'Glitch with {hits} hit{s}.'

    return result

def setup(bot):
    bot.add_cog(RPG(bot))
<commit_msg>Check for totaling correctly<commit_after>import asyncio
from concurrent import futures
import os
import random
import re
from urllib.parse import parse_qs

import discord
from discord.ext import commands
from lxml import etree

from starwars import starroller, die_names

class RPG:
    def __init__(self, bot):
        self.bot = bot
        self.tarot_url = 'http://www.tarotlore.com/tarot-cards/{}/'

    @commands.command()
    async def choose(self, ctx, *options):
        """Choose between some options. Use quotes if they have spaces."""
        await ctx.send(random.choice(options))

    @commands.command()
    async def tarot(self, ctx, *suits):
        """Get a random tarot card.

        You can specify the suits from which to pull, options are:
        minor:
            cups
            swords
            wands
            pentacles
        major"""
        async with ctx.typing():
            cards = []
            if not suits:
                suits = ('cups', 'swords', 'wands', 'pentacles', 'major')
            if 'minor' in suits:
                suits = suits + ('cups', 'swords', 'wands', 'pentacles')
            suits = set(suit.lower() for suit in suits)
            for root, dirs, files in os.walk('data/tarot'):
                if any(suit in root for suit in suits):
                    cards += [f'{root}/{card}' for card in files]
            try:
                card = random.choice(cards).replace('\\', '/')
            except IndexError:
                await ctx.send('Please specify a valid suit, or no suit.')
                return
            match = re.match(r'(?:\w+\/)+[IVX0_]*([\w_]+)\.jpg', card)
            name = match.groups()[0].replace('_', ' ')
            url = self.tarot_url.format(name.lower().replace(' ', '-'))
            embed = discord.Embed()
            embed.title = name
            embed.url = url
            filename = card.replace('_', '').rpartition('/')[-1]
            embed.set_image(url=f'attachment://{filename}')
            await ctx.send(file=discord.File(f'{card}', filename),
                           embed=embed)

    @commands.command(aliases=['r'])
    async def roll(self, ctx, *, inp='1d20'):
        """Roll some dice!

        Can roll multiple dice of any size, with modifiers.
        Format: XdY([+-^v]Z)(xN)(s)(t)
        X is the number of dice
        Y is the number of sides
        + adds Z to the result
        - subtracts Z from the result
        ^ drops the Z highest dice
        v drops the Z lowest dice
        x repeats the roll N times
        s sorts the results
        t totals each roll
        You can join rolls with commas
        """
        if inp == 'stats':
            inp = '4d6v1x6t'
        inp = ''.join(inp.split()).lower()
        expr = r'^([\d]*d?\d+([+\-^v]\d+)?(x\d+)?([ts]{1,2})?,?)+(?<!,)$'
        if re.match(expr, inp) is None:
            raise commands.BadArgument

        rolls = inp.split(',')
        args_batch = []
        for roll in rolls:
            if 'd' not in roll:
                roll = f'1d{roll}'
            elif roll[0] == 'd':
                roll = f'1{roll}'
            args = tuple(int(arg) for arg in re.findall(r'\d+', roll))

            num = args[0]
            sides = args[1]

            hi_drop = 0
            lo_drop = 0
            mod = 0

            if '^' in inp:
                hi_drop = args[2]
            elif 'v' in inp:
                lo_drop = args[2]
            elif '+' in inp:
                mod = args[2]
            elif '-' in inp:
                mod = -args[2]

            if 'x' in inp:
                times = args[-1]
            else:
                times = 1

            loop = self.bot.loop
            args = (num, sides, lo_drop, hi_drop, mod, times)
            args_batch.append(args)

        future = loop.run_in_executor(None, self.roll_helper, args_batch)
        async with ctx.typing():
            results = await asyncio.wait_for(future, 10, loop=loop)

        out = []
        for roll, result in zip(rolls, results):
            total = 't' in roll
            if total:
                result = [[sum(roll_)] for roll_ in result]
            if 's' in inp:
                for roll_ in result:
                    roll_.sort()
                result.sort()
            if total or len(result) == 1:
                result = [roll_[0] for roll_ in result]
            if 'x' not in inp:
                result = result[0]
            out.append(f'{roll}: {result}')
        await ctx.reply('\n'.join(out))

    def roll_helper(self, rolls):
        out = []
        for roll in rolls:
            out.append(roller(*roll))
        return out

    @roll.error
    async def roll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, (commands.MissingRequiredArgument,
                      commands.BadArgument)):
            await ctx.send('Invalid input. Valid input examples:'
                           '\n1d20+3'
                           '\n1d6'
                           '\n2d8-4'
                           '\n2d20^1'
                           '\n4d6v1x6t')
        elif isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        elif isinstance(e, discord.HTTPException):
            await ctx.reply('Your results were too long. Maybe sum them?')
        else:
            await self.bot.handle_error(ctx, e)

    @commands.command(aliases=['shadroll', 'sr'])
    async def shadowroll(self, ctx, *, inp):
        """Roll some dice - for Shadowrun!

        Format: N[e]
        Roll N six-sided dice and return the number of dice that rolled 5 or 6.
        If you put "e" after the number, 6s are counted and then rerolled."""
        inp = inp.strip()
        expr = r'^\d+e?$'
        if not re.match(expr, inp):
            raise commands.BadArgument
        edge = 'e' in inp
        if edge:
            inp = inp[:-1]
        num = int(inp)

        loop = self.bot.loop
        args = (num, edge)
        future = loop.run_in_executor(None, shadowroller, *args)
        async with ctx.typing():
            result = await asyncio.wait_for(future, 10, loop=loop)

        await ctx.reply(result)

    @shadowroll.error
    async def shadowroll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, (commands.MissingRequiredArgument,
                      commands.BadArgument)):
            await ctx.send('Invalid input. Valid input examples:'
                           '\n6'
                           '\n13e')
        elif isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        else:
            await self.bot.handle_error(ctx, e)

    @commands.command(aliases=['sw'])
    async def starroll(self, ctx, *, inp):
        """Roll some dice - for Fantasy Flight Star Wars!

        Available dice:
        b[oost]
        a[bility]
        p[roficiency]
        s[etback]
        d[ifficulty]
        c[hallenge]
        f[orce]

        Input examples:
        4a3d
        3a2p1b4d1c
        2f"""
        inp = inp.lower()
        expr = r'^(\d+[a-z])+$'
        match = re.match(expr, inp)
        if not match:
            raise commands.BadArgument
        dice = {}
        for roll in match.groups():
            num = int(roll[:-1])
            try:
                die = die_names[roll[-1]]
            except KeyError:
                await ctx.send(f'Die "{roll[-1]}" does not exist.')
                return
            dice[die] = num

        loop = self.bot.loop
        future = loop.run_in_executor(None, lambda: starroller(**dice))
        async with ctx.typing():
            try:
                result = await asyncio.wait_for(future, 10, loop=loop)
            except ValueError:
                await ctx.send('Force dice cannot be used with other dice.')
            else:
                await ctx.reply(result)

    @starroll.error
    async def starroll_error(self, ctx, e):
        e = getattr(e, 'original', e)
        if isinstance(e, futures.TimeoutError):
            await ctx.reply('Your execution took too long. Roll fewer dice.')
        else:
            await self.bot.handle_error(ctx, e)

def roller(num=1, sides=20, lo_drop=0, hi_drop=0, mod=0, times=1):
    rolls = []
    for _ in range(times):
        pool = [random.randint(1, sides) for _ in range(num)]
        if lo_drop or hi_drop:
            sorted_pool = sorted(pool)
            dropped_vals = sorted_pool[:lo_drop] + sorted_pool[num-hi_drop:]
            for val in dropped_vals:
                pool.remove(val)
        if mod:
            pool = [sum(pool) + mod]
        rolls.append(pool)
    return rolls

def shadowroller(num, edge=False):
    rolls = hits = count1 = 0
    while True:
        count6 = 0
        rolls += num
        for _ in range(num):
            roll = random.randint(1, 6)
            if roll > 4:
                hits += 1
                if roll == 6:
                    count6 += 1
            elif roll == 1:
                count1 += 1
        if not (count6 > 0 and edge):
            break
        num = count6
    s = 's' if hits != 1 else ''
    if count1 < rolls / 2:
        result = f'{hits} hit{s}.'
    elif hits == 0:
        result = 'Critical glitch.'
    else:
        result = f'Glitch with {hits} hit{s}.'

    return result

def setup(bot):
    bot.add_cog(RPG(bot))
<|endoftext|>


<commit_before>import feedparser
import operator
import hashlib
from urllib.request import urlopen
from sopel.tools import SopelMemory
from sopel.module import commands, interval, NOLIMIT, require_privmsg, require_admin
from sopel.config.types import (StaticSection, ListAttribute)

class RSSSection(StaticSection):
    del_by_id = ListAttribute('del_by_id', default=False)
    feeds = ListAttribute('feeds', default=[])
    hashes = ListAttribute('hashes', default=[])
    monitoring_channel = ListAttribute('monitoring_channel',default=[])

def setup(bot):
    bot.config.define_section('rss', RSSSection)
    bot.memory['rss'] = SopelMemory()
    bot.memory['rss']['del_by_id'] = False
    bot.memory['rss']['feeds'] = []
    bot.memory['rss']['hashes'] = RingBuffer(600)
    bot.memory['rss']['monitoring_channel'] = []
    __config_read(bot)

def configure(config):
    config.define_section('rss', RSSSection)
    config.rss.configure_setting('del_by_id', 'allow to delete feeds by index. this is potentially dangerous as the index of the remaining feeds may change after one feed has been deleted')
    config.rss.configure_setting('feeds', 'comma separated strings consisting of channel, name, url separated by spaces')
    config.rss.configure_setting('hashes', 'comma separated strings consisting of hashes of rss feed entries')
    config.rss.configure_setting('monitoring_channel', 'channel in which the bot will report errors')

def shutdown(bot):
    print('Shutting down!')
    bot.debug('rss', 'Shutting down!', 'always')
    __config_save(bot)

@require_admin
@commands('rssadd')
def rssadd(bot, trigger):
    # check number of parameters
    for i in range(2,6):
        if trigger.group(i) is None:
            bot.say('syntax: {}{} <channel> <name> <url>'.format(bot.config.core.prefix, trigger.group(1)))
            return NOLIMIT

    # get arguments
    channel = trigger.group(3)
    name = trigger.group(4)
    url = trigger.group(5)

    # check that feed name is not a number
    try:
        int(name)
    except ValueError:
        pass
    else:
        bot.say('feed name "{}" is invalid, it has to contain at least one letter'.format(name))
        return NOLIMIT

    # check that feed name is unique
    for feed in bot.memory['rss']['feeds']:
        if name == feed['name']:
            bot.say('feed name "{}" is already in use, please choose a different one'.format(name))
            return NOLIMIT

    # check that channel starts with #
    if not channel.startswith('#'):
        bot.say('channel "{}" must begin with a "#"'.format(channel))
        return NOLIMIT

    # check that feed url is online
    try:
        with urlopen(url) as f:
            if f.status == 200:
                # save config to memory
                bot.memory['rss']['feeds'].append({'channel':channel, 'name':name, 'url':url})
                bot.say('rss feed "{}" added successfully'.format(url))
                bot.join(channel)
    except:
        bot.say('unable to add feed "{}": invalid url! syntax: .{} <channel> <name> <url>'.format(url, trigger.group(1)))
        return NOLIMIT

    __config_save(bot)

    return NOLIMIT

@require_admin
@commands('rssdel')
def rssdel(bot, trigger):
    # get arguments
    arg = trigger.group(3)

    if arg is None:
        if bot.memory['rss']['del_by_id']:
            bot.say('syntax: {}{} <id>|<name>'.format(bot.config.core.prefix, trigger.group(1)))
        else:
            bot.say('syntax: {}{} <name>'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    # check if arg is a number
    try:
        int(arg)
    except ValueError:
        # arg is not a number: name
        __delFeedByName(bot, arg)
    else:
        # arg is a number: id
        if not bot.memory['rss']['del_by_id']:
            # subtract 1 from id argument as rsslist starts counting at 1 and not at 0
            bot.say('please use ".{} {}" to delete feed with the current id {}'.format(trigger.group(1), __getNameByIndex(bot, int(arg) - 1), arg))
            return NOLIMIT

        # subtract 1 from id argument as rsslist starts counting at 1 and not at 0
        __delFeedByIndex(bot, int(arg) - 1)

    __config_save(bot)

    return NOLIMIT

# shoot yourself in the feet
@require_privmsg
@require_admin
@commands('rssdelallfeeds')
def rssdeleteallfeeds(bot, trigger):
    # check number of parameters
    if not trigger.group(2) is None:
        bot.say('syntax: {}{}'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    bot.memory['rss']['feeds'].clear()
    bot.say('all rss feeds deleted successfully')

    __config_save(bot)

    return NOLIMIT

@require_admin
@commands('rssget')
def rssget(bot, trigger):
    # check number of parameters
    if trigger.group(3) is None or (trigger.group(4) and trigger.group(4) != 'all') or trigger.group(5):
        bot.say('syntax: {}{} <name> [<scope>]'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    # if chatty is True then the bot will post feed items to a channel
    chatty = False

    name = trigger.group(3)

    if (trigger.group(4) and trigger.group(4) == 'all'):
        chatty = True

    __updateFeed(bot, name, chatty)

    return NOLIMIT

@require_admin
@commands('rssjoin')
def rssjoin(bot, trigger):
    # check number of parameters
    if not trigger.group(2) is None:
        bot.say('syntax: {}{}'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    for feed in bot.memory['rss']['feeds']:
        bot.join(feed['channel'])

    if bot.memory['rss']['monitoring_channel']:
        bot.join(bot.memory['rss']['monitoring_channel'])

    __config_save(bot)

@require_admin
@commands('rsslist')
def rsslist(bot, trigger):
    # check number of parameters
    if not trigger.group(4) is None:
        bot.say('syntax: {}{} [<channel>]'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    channel = trigger.group(3)

    # check that channel starts with #
    if channel and not channel.startswith('#'):
        bot.say('channel "{}" must begin with a "#"'.format(channel))
        return NOLIMIT

    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if channel:
            if channel != feed['channel']:
                continue
        bot.say('{}: {} {} {}'.format(feeds.index(feed) + 1, feed['channel'], feed['name'], feed['url']))
    return NOLIMIT

# read config from disk to memory
def __config_read(bot):
    # feeds
    if bot.config.rss.feeds and bot.config.rss.feeds[0]:
        for feed in bot.config.rss.feeds:
            # split feed line by spaces
            atoms = feed.split(' ')

            bot.memory['rss']['feeds'].append({'channel': atoms[0], 'name': atoms[1], 'url': atoms[2]})

    # del_by_id
    if bot.config.rss.del_by_id:
        bot.memory['rss']['del_by_id'] = bot.config.rss.del_by_id

    # monitoring_channel
    if bot.config.rss.monitoring_channel:
        bot.memory['rss']['moitoring_channel'] = bot.config.rss.monitoring_channel

    # sort list by channel name
    bot.memory['rss']['feeds'].sort(key=operator.itemgetter('channel'))

    # read hashes
    for hash in bot.config.rss.hashes:
        bot.memory['rss']['hashes'].append(hash)

    # write config to disk after it has been sorted
    __config_save(bot)

# save config from memory to disk
def __config_save(bot):

    if not bot.memory['rss']['feeds']:
        return NOLIMIT

    # monitoring_channel
    bot.config.rss.monitoring_channel = bot.memory['rss']['monitoring_channel']

    # sort list by channel name
    bot.memory['rss']['feeds'].sort(key=operator.itemgetter('channel'))

    # flatten feeds for config file
    feeds = []
    for feed in bot.memory['rss']['feeds']:
        feeds.append(feed['channel'] + ' ' + feed['name'] + ' ' + feed['url'])
    bot.config.rss.feeds = [",".join(feeds)]

    # save channels
    channels = bot.config.core.channels
    for feed in bot.memory['rss']['feeds']:
        if not feed['channel'] in channels:
            bot.config.core.channels += [feed['channel']]

    # save hashes
    bot.config.rss.hashes = []
    for hash in bot.memory['rss']['hashes'].get():
        bot.config.rss.hashes += [hash]

    try:
        bot.config.save()
    except:
        print('Unable to save config to disk!')
        bot.debug('rss', 'Unable to save config to disk!', 'always')

def __delFeedByIndex(bot, index):
    # read url of feed in order to give feedback after the feed has been deleted
    try:
        url = __getUrlByIndex(bot, index)
    except:
        bot.say('feed number "{}" doesn\'t exist!'.format(index))
        return NOLIMIT

    # delete feed
    try:
        if bot.memory['rss']['feeds'][index]:
            del((bot.memory['rss']['feeds'])[index])
            bot.say('rss feed "{}" deleted successfully'.format(url))
    except:
        bot.say('unable to delete feed "{}": feed number doesn\'t exist!'.format(index))

def __delFeedByName(bot, name):
    index = __getIndexByName(bot, name)
    __delFeedByIndex(bot, index)

def __getChannelByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feed['channel']

def __getUrlByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feed['url']

def __getIndexByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feeds.index(feed)

def __getNameByIndex(bot, index):
    return bot.memory['rss']['feeds'][index]['name']

def __getUrlByIndex(bot, index):
    return bot.memory['rss']['feeds'][index]['url']

def __hashEntry(entry):
    return hashlib.md5(entry.encode('utf-8')).hexdigest()

@interval(60)
def __update(bot):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        __updateFeed(bot, feed['name'], False)

def __updateFeed(bot, name, chatty):
    try:
        url = __getUrlByName(bot, name)
        if url is None:
            raise Exception
    except:
        bot.say('url of feed "{}" couldn\'t be read!'.format(name), bot.memory['rss']['monitoring_channel'])
        return NOLIMIT

    try:
        feed = feedparser.parse(url)
    except:
        bot.say('error reading feed "{}"'.format(url), bot.memory['rss']['monitoring_channel'])
        return NOLIMIT

    channel = __getChannelByName(bot, name)

    # print new or all items
    for item in reversed(feed['entries']):
        hash = __hashEntry(name + item['title'] + item['link'] + item['summary'])
        if chatty or not hash in bot.memory['rss']['hashes'].get():
            bot.memory['rss']['hashes'].append(hash)
            message = '\u0002[' + name + ']\u000F '
            message += item['title'] + ' \u0002→\u000F ' + item['link']
            bot.say(message, channel)

    # write config to disk after new hashes have been calculated
    __config_save(bot)

# Implementing a Ring Buffer
# https://www.safaribooksonline.com/library/view/python-cookbook/0596001673/ch05s19.html
class RingBuffer:
    """ class that implements a not-yet-full buffer """
    def __init__(self,size_max):
        self.max = size_max
        self.index = 0
        self.data = []

    class __Full:
        """ class that implements a full buffer """
        def append(self, x):
            """ Append an element overwriting the oldest one. """
            self.data[self.cur] = x
            self.cur = (self.cur+1) % self.max
        def get(self):
            """ return list of elements in correct order """
            return self.data[self.cur:]+self.data[:self.cur]

    def append(self,x):
        """ append an element at the end of the buffer """
        self.data.append(x)
        if len(self.data) == self.max:
            self.cur = 0
            # Permanently change self's class from non-full to full
            self.__class__ = self.__Full

    def get(self):
        """ return a list of elements from the oldest to the newest. """
        return self.data<commit_msg>Enlarge RingBuffer and fix typo<commit_after>import feedparser
import operator
import hashlib
from urllib.request import urlopen
from sopel.tools import SopelMemory
from sopel.module import commands, interval, NOLIMIT, require_privmsg, require_admin
from sopel.config.types import (StaticSection, ListAttribute)

class RSSSection(StaticSection):
    del_by_id = ListAttribute('del_by_id', default=False)
    feeds = ListAttribute('feeds', default=[])
    hashes = ListAttribute('hashes', default=[])
    monitoring_channel = ListAttribute('monitoring_channel',default=[])

def setup(bot):
    bot.config.define_section('rss', RSSSection)
    bot.memory['rss'] = SopelMemory()
    bot.memory['rss']['del_by_id'] = False
    bot.memory['rss']['feeds'] = []
    bot.memory['rss']['hashes'] = RingBuffer(1000)
    bot.memory['rss']['monitoring_channel'] = []
    __config_read(bot)

def configure(config):
    config.define_section('rss', RSSSection)
    config.rss.configure_setting('del_by_id', 'allow to delete feeds by index. this is potentially dangerous as the index of the remaining feeds may change after one feed has been deleted')
    config.rss.configure_setting('feeds', 'comma separated strings consisting of channel, name, url separated by spaces')
    config.rss.configure_setting('hashes', 'comma separated strings consisting of hashes of rss feed entries')
    config.rss.configure_setting('monitoring_channel', 'channel in which the bot will report errors')

def shutdown(bot):
    print('Shutting down!')
    bot.debug('rss', 'Shutting down!', 'always')
    __config_save(bot)

@require_admin
@commands('rssadd')
def rssadd(bot, trigger):
    # check number of parameters
    for i in range(2,6):
        if trigger.group(i) is None:
            bot.say('syntax: {}{} <channel> <name> <url>'.format(bot.config.core.prefix, trigger.group(1)))
            return NOLIMIT

    # get arguments
    channel = trigger.group(3)
    name = trigger.group(4)
    url = trigger.group(5)

    # check that feed name is not a number
    try:
        int(name)
    except ValueError:
        pass
    else:
        bot.say('feed name "{}" is invalid, it has to contain at least one letter'.format(name))
        return NOLIMIT

    # check that feed name is unique
    for feed in bot.memory['rss']['feeds']:
        if name == feed['name']:
            bot.say('feed name "{}" is already in use, please choose a different one'.format(name))
            return NOLIMIT

    # check that channel starts with #
    if not channel.startswith('#'):
        bot.say('channel "{}" must begin with a "#"'.format(channel))
        return NOLIMIT

    # check that feed url is online
    try:
        with urlopen(url) as f:
            if f.status == 200:
                # save config to memory
                bot.memory['rss']['feeds'].append({'channel':channel, 'name':name, 'url':url})
                bot.say('rss feed "{}" added successfully'.format(url))
                bot.join(channel)
    except:
        bot.say('unable to add feed "{}": invalid url! syntax: .{} <channel> <name> <url>'.format(url, trigger.group(1)))
        return NOLIMIT

    __config_save(bot)

    return NOLIMIT

@require_admin
@commands('rssdel')
def rssdel(bot, trigger):
    # get arguments
    arg = trigger.group(3)

    if arg is None:
        if bot.memory['rss']['del_by_id']:
            bot.say('syntax: {}{} <id>|<name>'.format(bot.config.core.prefix, trigger.group(1)))
        else:
            bot.say('syntax: {}{} <name>'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    # check if arg is a number
    try:
        int(arg)
    except ValueError:
        # arg is not a number: name
        __delFeedByName(bot, arg)
    else:
        # arg is a number: id
        if not bot.memory['rss']['del_by_id']:
            # subtract 1 from id argument as rsslist starts counting at 1 and not at 0
            bot.say('please use ".{} {}" to delete feed with the current id {}'.format(trigger.group(1), __getNameByIndex(bot, int(arg) - 1), arg))
            return NOLIMIT

        # subtract 1 from id argument as rsslist starts counting at 1 and not at 0
        __delFeedByIndex(bot, int(arg) - 1)

    __config_save(bot)

    return NOLIMIT

# shoot yourself in the feet
@require_privmsg
@require_admin
@commands('rssdelallfeeds')
def rssdeleteallfeeds(bot, trigger):
    # check number of parameters
    if not trigger.group(2) is None:
        bot.say('syntax: {}{}'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    bot.memory['rss']['feeds'].clear()
    bot.say('all rss feeds deleted successfully')

    __config_save(bot)

    return NOLIMIT

@require_admin
@commands('rssget')
def rssget(bot, trigger):
    # check number of parameters
    if trigger.group(3) is None or (trigger.group(4) and trigger.group(4) != 'all') or trigger.group(5):
        bot.say('syntax: {}{} <name> [<scope>]'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    # if chatty is True then the bot will post feed items to a channel
    chatty = False

    name = trigger.group(3)

    if (trigger.group(4) and trigger.group(4) == 'all'):
        chatty = True

    __updateFeed(bot, name, chatty)

    return NOLIMIT

@require_admin
@commands('rssjoin')
def rssjoin(bot, trigger):
    # check number of parameters
    if not trigger.group(2) is None:
        bot.say('syntax: {}{}'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    for feed in bot.memory['rss']['feeds']:
        bot.join(feed['channel'])

    if bot.memory['rss']['monitoring_channel']:
        bot.join(bot.memory['rss']['monitoring_channel'])

    __config_save(bot)

@require_admin
@commands('rsslist')
def rsslist(bot, trigger):
    # check number of parameters
    if not trigger.group(4) is None:
        bot.say('syntax: {}{} [<channel>]'.format(bot.config.core.prefix, trigger.group(1)))
        return NOLIMIT

    channel = trigger.group(3)

    # check that channel starts with #
    if channel and not channel.startswith('#'):
        bot.say('channel "{}" must begin with a "#"'.format(channel))
        return NOLIMIT

    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if channel:
            if channel != feed['channel']:
                continue
        bot.say('{}: {} {} {}'.format(feeds.index(feed) + 1, feed['channel'], feed['name'], feed['url']))
    return NOLIMIT

# read config from disk to memory
def __config_read(bot):
    # feeds
    if bot.config.rss.feeds and bot.config.rss.feeds[0]:
        for feed in bot.config.rss.feeds:
            # split feed line by spaces
            atoms = feed.split(' ')
            bot.memory['rss']['feeds'].append({'channel': atoms[0], 'name': atoms[1], 'url': atoms[2]})

    # del_by_id
    if bot.config.rss.del_by_id:
        bot.memory['rss']['del_by_id'] = bot.config.rss.del_by_id

    # monitoring_channel
    if bot.config.rss.monitoring_channel:
        bot.memory['rss']['monitoring_channel'] = bot.config.rss.monitoring_channel

    # sort list by channel name
    bot.memory['rss']['feeds'].sort(key=operator.itemgetter('channel'))

    # read hashes
    for hash in bot.config.rss.hashes:
        bot.memory['rss']['hashes'].append(hash)

    # write config to disk after it has been sorted
    __config_save(bot)

# save config from memory to disk
def __config_save(bot):

    if not bot.memory['rss']['feeds']:
        return NOLIMIT

    # monitoring_channel
    bot.config.rss.monitoring_channel = bot.memory['rss']['monitoring_channel']

    # sort list by channel name
    bot.memory['rss']['feeds'].sort(key=operator.itemgetter('channel'))

    # flatten feeds for config file
    feeds = []
    for feed in bot.memory['rss']['feeds']:
        feeds.append(feed['channel'] + ' ' + feed['name'] + ' ' + feed['url'])
    bot.config.rss.feeds = [",".join(feeds)]

    # save channels
    channels = bot.config.core.channels
    for feed in bot.memory['rss']['feeds']:
        if not feed['channel'] in channels:
            bot.config.core.channels += [feed['channel']]

    # save hashes
    bot.config.rss.hashes = []
    for hash in bot.memory['rss']['hashes'].get():
        bot.config.rss.hashes += [hash]

    try:
        bot.config.save()
    except:
        print('Unable to save config to disk!')
        bot.debug('rss', 'Unable to save config to disk!', 'always')

def __delFeedByIndex(bot, index):
    # read url of feed in order to give feedback after the feed has been deleted
    try:
        url = __getUrlByIndex(bot, index)
    except:
        bot.say('feed number "{}" doesn\'t exist!'.format(index))
        return NOLIMIT

    # delete feed
    try:
        if bot.memory['rss']['feeds'][index]:
            del((bot.memory['rss']['feeds'])[index])
            bot.say('rss feed "{}" deleted successfully'.format(url))
    except:
        bot.say('unable to delete feed "{}": feed number doesn\'t exist!'.format(index))

def __delFeedByName(bot, name):
    index = __getIndexByName(bot, name)
    __delFeedByIndex(bot, index)

def __getChannelByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feed['channel']

def __getUrlByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feed['url']

def __getIndexByName(bot, name):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        if feed['name'] == name:
            return feeds.index(feed)

def __getNameByIndex(bot, index):
    return bot.memory['rss']['feeds'][index]['name']

def __getUrlByIndex(bot, index):
    return bot.memory['rss']['feeds'][index]['url']

def __hashEntry(entry):
    return hashlib.md5(entry.encode('utf-8')).hexdigest()

@interval(60)
def __update(bot):
    feeds = bot.memory['rss']['feeds']
    for feed in feeds:
        __updateFeed(bot, feed['name'], False)

def __updateFeed(bot, name, chatty):
    try:
        url = __getUrlByName(bot, name)
        if url is None:
            raise Exception
    except:
        bot.say('url of feed "{}" couldn\'t be read!'.format(name), bot.memory['rss']['monitoring_channel'])
        return NOLIMIT

    try:
        feed = feedparser.parse(url)
    except:
        bot.say('error reading feed "{}"'.format(url), bot.memory['rss']['monitoring_channel'])
        return NOLIMIT

    channel = __getChannelByName(bot, name)

    # print new or all items
    for item in reversed(feed['entries']):
        hash = __hashEntry(name + item['title'] + item['link'] + item['summary'])
        if chatty or not hash in bot.memory['rss']['hashes'].get():
            bot.memory['rss']['hashes'].append(hash)
            message = '\u0002[' + name + ']\u000F '
            message += item['title'] + ' \u0002→\u000F ' + item['link']
            bot.say(message, channel)

    # write config to disk after new hashes have been calculated
    __config_save(bot)

# Implementing a Ring Buffer
# https://www.safaribooksonline.com/library/view/python-cookbook/0596001673/ch05s19.html
class RingBuffer:
    """ class that implements a not-yet-full buffer """
    def __init__(self,size_max):
        self.max = size_max
        self.index = 0
        self.data = []

    class __Full:
        """ class that implements a full buffer """
        def append(self, x):
            """ Append an element overwriting the oldest one. """
            self.data[self.cur] = x
            self.cur = (self.cur+1) % self.max
        def get(self):
            """ return list of elements in correct order """
            return self.data[self.cur:]+self.data[:self.cur]

    def append(self,x):
        """ append an element at the end of the buffer """
        self.data.append(x)
        if len(self.data) == self.max:
            self.cur = 0
            # Permanently change self's class from non-full to full
            self.__class__ = self.__Full

    def get(self):
        """ return a list of elements from the oldest to the newest. """
        return self.data<|endoftext|>


<commit_before>from datetime import datetime
import time
from discord.ext import commands
from discord import utils
import aiosqlite3
import os

class ConnectionHolder:
    def __init__(self):
        self.dbcon = None

    async def load_db(self, database_name, loop):
        if not os.path.isfile(database_name):
            print(f'Creating database {database_name}')
            with open('schema.sql', 'r', encoding='utf-8') as f:
                schema = f.read()
            self.dbcon = await aiosqlite3.connect(database_name, loop=loop)
            await self.dbcon.executescript(schema)
            await self.dbcon.commit()
            print(f'{database_name} initialized')
        else:
            self.dbcon = await aiosqlite3.connect(database_name, loop=loop)
            print(f'Loaded {database_name}')

    async def __aenter__(self):
        self.dbcon.__enter__()
        cursor = await self.dbcon.cursor()
        return cursor

    async def __aexit__(self, exc_class, exc, traceback):
        self.dbcon.__exit__(exc_class, exc, traceback)
        if self.dbcon.in_transaction:
            await self.dbcon.commit()

class DatabaseCog(commands.Cog):
    """
    Base class for cogs that interact with the database.
    """
    def __init__(self, bot):
        self.bot = bot
        print(f'Cog "{self.qualified_name}" loaded')

    async def add_restriction(self, user_id, role):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM permanent_roles WHERE user_id=? AND role_id=?', (user_id, role.id))
            if await cur.fetchone() is None:
                await cur.execute('INSERT INTO permanent_roles VALUES(?, ?)', (user_id, role.id))
                return True
            return False

    async def remove_restriction(self, user_id, role):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM permanent_roles WHERE user_id=? AND role_id=?', (user_id, role.id))

    async def get_restrictions_roles_id(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT role_id FROM permanent_roles WHERE user_id=?', (user_id,))
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def add_staff(self, user_id, position):
        async with self.bot.holder as cur:
            if await self.get_stafftrole(user_id) is None:
                await cur.execute('INSERT INTO staff VALUES(?, ?)', (user_id, position))
            else:
                await cur.execute('UPDATE staff SET position=? WHERE user_id=? ', (position, user_id))

    async def remove_staff(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM staff WHERE user_id=?', (user_id,))

    async def get_staff(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM staff')
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def get_helpers(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM helpers')
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def add_helper(self, user_id, console):
        async with self.bot.holder as cur:
            if await self.get_console(user_id) is None:
                await cur.execute('INSERT INTO helpers VALUES(?, ?)', (user_id, console))
            else:
                await cur.execute('UPDATE helpers SET console=? WHERE user_id=?', (console, user_id))

    async def remove_helper(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM helpers WHERE user_id=?', (user_id,))

    async def get_console(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT console FROM helpers WHERE user_id=?', (user_id,))
            row = await cur.fetchone()
            return row[0] if row is not None else row

    async def get_stafftrole(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT position FROM staff WHERE user_id=?', (user_id,))
            rank = await cur.fetchone()
            return rank[0] if rank is not None else rank

    async def add_warn(self, user_id, issuer_id, reason):
        async with self.bot.holder as cur:
            snowflake = utils.time_snowflake(datetime.now())
            await cur.execute('INSERT INTO warns VALUES(?, ?, ?, ?)', (snowflake, user_id, issuer_id, reason))

    async def remove_warn_id(self, user_id, index):
        # i dont feel so good
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM warns WHERE warn_id in (SELECT warn_id FROM warns WHERE user_id=? LIMIT ?,1)', (user_id, index-1))

    async def remove_warns(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM warns WHERE user_id=?', (user_id,))

    async def get_warns(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM warns WHERE user_id=?', (user_id,))
            return await cur.fetchall()

    async def add_timed_restriction(self, user_id, end_date, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT 1 FROM timed_restrictions WHERE user_id=? AND type=?', (user_id, type))
            if await cur.fetchone() is not None:
                await cur.execute('UPDATE timed_restrictions SET timestamp=?, alert=0 WHERE user_id=? AND type=?', (end_date, user_id, type))
            else:
                await cur.execute('INSERT INTO timed_restrictions VALUES(?, ?, ?, ?)', (user_id, end_date, type, 0))

    async def remove_timed_restriction(self, user_id, type):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM timed_restrictions WHERE user_id=? AND type=?', (user_id, type))

    async def get_time_restrictions_by_user_type(self, userid, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * from timed_restrictions WHERE user_id=? AND type=?', (userid,type))
            return await cur.fetchone()

    async def get_time_restrictions_by_type(self, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * from timed_restrictions WHERE type=?', (type,))
            return await cur.fetchall()

    async def set_time_restriction_alert(self, user_id, type):
        async with self.bot.holder as cur:
            await cur.execute('UPDATE timed_restrictions SET alert=1 WHERE user_id=? AND type=?', (user_id, type))

    async def add_softban(self, user_id, issuer_id, reason, timestamp=None):
        if not timestamp:
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO softbans VALUES(?, ? , ?, ?)', (user_id, issuer_id, reason, timestamp))

    async def remove_softban(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM softbans WHERE user_id = ?', (user_id,))

    async def get_softban(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM softbans WHERE user_id=?', (user_id,))
            return await cur.fetchone()

    async def add_watch(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO watchlist VALUES(?)', (user_id,))

    async def remove_watch(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM watchlist WHERE user_id=?', (user_id,))

    async def is_watched(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM watchlist WHERE user_id=?', (user_id,))
            return await cur.fetchone() is not None

    async def add_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO nofilter VALUES(?)', (channel.id,))

    async def remove_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM nofilter WHERE channel_id=?', (channel.id,))

    async def check_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('SELECT channel_id FROM nofilter WHERE channel_id=?', (channel.id,))
            return await cur.fetchone() is not None

    async def add_friendcode(self, user_id, fc):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO friend_codes VALUES (?,?)', (user_id, fc))

    async def get_friendcode(self, user_id):
        async with self.bot.holder as cur:
            return await cur.execute('SELECT * FROM friend_codes WHERE user_id = ?', (user_id,))

    async def delete_friendcode(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM friend_codes WHERE user_id = ?', (user_id,))
<commit_msg>Add get_staff_role and get_helpers_role method<commit_after>from datetime import datetime
import time
from discord.ext import commands
from discord import utils
import aiosqlite3
import os

class ConnectionHolder:
    def __init__(self):
        self.dbcon = None

    async def load_db(self, database_name, loop):
        if not os.path.isfile(database_name):
            print(f'Creating database {database_name}')
            with open('schema.sql', 'r', encoding='utf-8') as f:
                schema = f.read()
            self.dbcon = await aiosqlite3.connect(database_name, loop=loop)
            await self.dbcon.executescript(schema)
            await self.dbcon.commit()
            print(f'{database_name} initialized')
        else:
            self.dbcon = await aiosqlite3.connect(database_name, loop=loop)
            print(f'Loaded {database_name}')

    async def __aenter__(self):
        self.dbcon.__enter__()
        cursor = await self.dbcon.cursor()
        return cursor

    async def __aexit__(self, exc_class, exc, traceback):
        self.dbcon.__exit__(exc_class, exc, traceback)
        if self.dbcon.in_transaction:
            await self.dbcon.commit()

class DatabaseCog(commands.Cog):
    """
    Base class for cogs that interact with the database.
    """
    def __init__(self, bot):
        self.bot = bot
        print(f'Cog "{self.qualified_name}" loaded')

    async def add_restriction(self, user_id, role):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM permanent_roles WHERE user_id=? AND role_id=?', (user_id, role.id))
            if await cur.fetchone() is None:
                await cur.execute('INSERT INTO permanent_roles VALUES(?, ?)', (user_id, role.id))
                return True
            return False

    async def remove_restriction(self, user_id, role):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM permanent_roles WHERE user_id=? AND role_id=?', (user_id, role.id))

    async def get_restrictions_roles_id(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT role_id FROM permanent_roles WHERE user_id=?', (user_id,))
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def add_staff(self, user_id, position):
        async with self.bot.holder as cur:
            if await self.get_stafftrole(user_id) is None:
                await cur.execute('INSERT INTO staff VALUES(?, ?)', (user_id, position))
            else:
                await cur.execute('UPDATE staff SET position=? WHERE user_id=? ', (position, user_id))

    async def remove_staff(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM staff WHERE user_id=?', (user_id,))

    async def get_staff(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM staff')
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def get_staff_role(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM staff')
            rows = await cur.fetchall()
            return rows

    async def get_helpers(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM helpers')
            rows = await cur.fetchall()
            if rows:
                return [x[0] for x in rows]
            return []

    async def get_helpers_role(self):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM helpers')
            rows = await cur.fetchall()
            return rows

    async def add_helper(self, user_id, console):
        async with self.bot.holder as cur:
            if await self.get_console(user_id) is None:
                await cur.execute('INSERT INTO helpers VALUES(?, ?)', (user_id, console))
            else:
                await cur.execute('UPDATE helpers SET console=? WHERE user_id=?', (console, user_id))

    async def remove_helper(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM helpers WHERE user_id=?', (user_id,))

    async def get_console(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT console FROM helpers WHERE user_id=?', (user_id,))
            row = await cur.fetchone()
            return row[0] if row is not None else row

    async def get_stafftrole(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT position FROM staff WHERE user_id=?', (user_id,))
            rank = await cur.fetchone()
            return rank[0] if rank is not None else rank

    async def add_warn(self, user_id, issuer_id, reason):
        async with self.bot.holder as cur:
            snowflake = utils.time_snowflake(datetime.now())
            await cur.execute('INSERT INTO warns VALUES(?, ?, ?, ?)', (snowflake, user_id, issuer_id, reason))

    async def remove_warn_id(self, user_id, index):
        # i dont feel so good
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM warns WHERE warn_id in (SELECT warn_id FROM warns WHERE user_id=? LIMIT ?,1)', (user_id, index-1))

    async def remove_warns(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM warns WHERE user_id=?', (user_id,))

    async def get_warns(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM warns WHERE user_id=?', (user_id,))
            return await cur.fetchall()

    async def add_timed_restriction(self, user_id, end_date, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT 1 FROM timed_restrictions WHERE user_id=? AND type=?', (user_id, type))
            if await cur.fetchone() is not None:
                await cur.execute('UPDATE timed_restrictions SET timestamp=?, alert=0 WHERE user_id=? AND type=?', (end_date, user_id, type))
            else:
                await cur.execute('INSERT INTO timed_restrictions VALUES(?, ?, ?, ?)', (user_id, end_date, type, 0))

    async def remove_timed_restriction(self, user_id, type):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM timed_restrictions WHERE user_id=? AND type=?', (user_id, type))

    async def get_time_restrictions_by_user_type(self, userid, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * from timed_restrictions WHERE user_id=? AND type=?', (userid,type))
            return await cur.fetchone()

    async def get_time_restrictions_by_type(self, type):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * from timed_restrictions WHERE type=?', (type,))
            return await cur.fetchall()

    async def set_time_restriction_alert(self, user_id, type):
        async with self.bot.holder as cur:
            await cur.execute('UPDATE timed_restrictions SET alert=1 WHERE user_id=? AND type=?', (user_id, type))

    async def add_softban(self, user_id, issuer_id, reason, timestamp=None):
        if not timestamp:
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO softbans VALUES(?, ? , ?, ?)', (user_id, issuer_id, reason, timestamp))

    async def remove_softban(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM softbans WHERE user_id = ?', (user_id,))

    async def get_softban(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT * FROM softbans WHERE user_id=?', (user_id,))
            return await cur.fetchone()

    async def add_watch(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO watchlist VALUES(?)', (user_id,))

    async def remove_watch(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM watchlist WHERE user_id=?', (user_id,))

    async def is_watched(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('SELECT user_id FROM watchlist WHERE user_id=?', (user_id,))
            return await cur.fetchone() is not None

    async def add_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO nofilter VALUES(?)', (channel.id,))

    async def remove_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM nofilter WHERE channel_id=?', (channel.id,))

    async def check_nofilter(self, channel):
        async with self.bot.holder as cur:
            await cur.execute('SELECT channel_id FROM nofilter WHERE channel_id=?', (channel.id,))
            return await cur.fetchone() is not None

    async def add_friendcode(self, user_id, fc):
        async with self.bot.holder as cur:
            await cur.execute('INSERT INTO friend_codes VALUES (?,?)', (user_id, fc))

    async def get_friendcode(self, user_id):
        async with self.bot.holder as cur:
            return await cur.execute('SELECT * FROM friend_codes WHERE user_id = ?', (user_id,))

    async def delete_friendcode(self, user_id):
        async with self.bot.holder as cur:
            await cur.execute('DELETE FROM friend_codes WHERE user_id = ?', (user_id,))
<|endoftext|>


<commit_before>from .utils import checks, db, time
from discord.ext import commands
import discord
import asyncio
import asyncpg
import datetime

class Reminders(db.Table):
    id = db.PrimaryKeyColumn()

    expires = db.Column(db.Datetime, index=True)
    created = db.Column(db.Datetime, default="now() at time zone 'utc'")
    event = db.Column(db.String)
    extra = db.Column(db.JSON, default="'{}'::jsonb")

class Timer:
    __slots__ = ('args', 'kwargs', 'event', 'id', 'created_at', 'expires')

    def __init__(self, *, record):
        self.id = record['id']

        extra = record['extra']
        self.args = extra.get('args', [])
        self.kwargs = extra.get('kwargs', {})
        self.event = record['event']
        self.created_at = record['created']
        self.expires = record['expires']

    @classmethod
    def temporary(cls, *, expires, created, event, args, kwargs):
        pseudo = {
            'id': None,
            'extra': { 'args': args, 'kwargs': kwargs },
            'event': event,
            'created': created,
            'expires': expires
        }
        return cls(record=pseudo)

    def __eq__(self, other):
        try:
            return self.id == other.id
        except AttributeError:
            return False

    def __hash__(self):
        return hash(self.id)

    @property
    def human_delta(self):
        return time.human_timedelta(self.created_at)

    def __repr__(self):
        return f'<Timer created={self.created_at} expires={self.expires} event={self.event}>'

class Reminder:
    """Reminders to do something."""

    def __init__(self, bot):
        self.bot = bot
        self._have_data = asyncio.Event(loop=bot.loop)
        self._current_timer = None
        self._task = bot.loop.create_task(self.dispatch_timers())

    def __unload(self):
        self._task.cancel()

    async def __error(self, ctx, error):
        if isinstance(error, commands.BadArgument):
            await ctx.send(error)

    async def get_active_timers(self, *, connection=None, days=7):
        query = "SELECT * FROM reminders WHERE expires < (CURRENT_DATE + $1::interval) ORDER BY expires;"
        con = connection or self.bot.pool

        records = await con.fetch(query, datetime.timedelta(days=days))
        return [Timer(record=a) for a in records]

    async def wait_for_active_timers(self, *, connection=None, days=7):
        async with db.MaybeAcquire(connection=connection, pool=self.bot.pool) as con:
            timers = await self.get_active_timers(connection=con, days=days)
            if len(timers):
                self._have_data.set()
                return timers

            self._have_data.clear()
            self._current_timer = None
            await self._have_data.wait()
            return await self.get_active_timers(connection=con, days=days)

    async def call_timer(self, timer):
        # delete the timer
        query = "DELETE FROM reminders WHERE id=$1;"
        await self.bot.pool.execute(query, timer.id)

        # dispatch the event
        event_name = f'{timer.event}_timer_complete'
        self.bot.dispatch(event_name, timer)

    async def dispatch_timers(self):
        try:
            while not self.bot.is_closed():
                # can only asyncio.sleep for up to ~48 days reliably
                # so we're gonna cap it off at 40 days
                # see: http://bugs.python.org/issue20493
                timers = await self.wait_for_active_timers(days=40)
                timer = self._current_timer = timers[0]
                now = datetime.datetime.utcnow()

                if timer.expires >= now:
                    to_sleep = (timer.expires - now).total_seconds()
                    await asyncio.sleep(to_sleep)

                await self.call_timer(timer)
        except asyncio.CancelledError:
            pass
        except (OSError, discord.ConnectionClosed, asyncpg.PostgresConnectionError):
            self._task.cancel()
            self._task = self.bot.loop.create_task(self.dispatch_timers())

    async def short_timer_optimisation(self, seconds, timer):
        await asyncio.sleep(seconds)
        event_name = f'{timer.event}_timer_complete'
        self.bot.dispatch(event_name, timer)

    async def create_timer(self, *args, **kwargs):
        """Creates a timer.

        Parameters
        -----------
        when: datetime.datetime
            When the timer should fire.
        event: str
            The name of the event to trigger.
            Will transform to 'on_{event}_timer_complete'.
        \*args
            Arguments to pass to the event
        \*\*kwargs
            Keyword arguments to pass to the event
        connection: asyncpg.Connection
            Special keyword-only argument to use a specific connection
            for the DB request.

        Note
        ------
        Arguments and keyword arguments must be JSON serialisable.

        Returns
        --------
        :class:`Timer`
        """
        when, event, *args = args

        try:
            connection = kwargs.pop('connection')
        except KeyError:
            connection = self.bot.pool

        now = datetime.datetime.utcnow()
        timer = Timer.temporary(event=event, args=args, kwargs=kwargs, expires=when, created=now)
        delta = (when - now).total_seconds()
        if delta <= 60:
            # a shortcut for small timers
            self.bot.loop.create_task(self.short_timer_optimisation(delta, timer))
            return timer

        query = """INSERT INTO reminders (event, extra, expires)
                   VALUES ($1, $2::jsonb, $3)
                   RETURNING id;
                """

        row = await connection.fetchrow(query, event, { 'args': args, 'kwargs': kwargs }, when)
        timer.id = row[0]

        self._have_data.set()

        # check if this timer is earlier than our currently run timer
        if self._current_timer and when < self._current_timer.expires:
            # cancel the task and re-run it
            self._task.cancel()
            self._task = self.bot.loop.create_task(self.dispatch_timers())

        return timer

    @commands.command(aliases=['timer'])
    async def reminder(self, ctx, when: time.FutureTime, *, message: commands.clean_content = 'something'):
        """Reminds you of something after a certain amount of time.

        The time can be any direct date (e.g. YYYY-MM-DD) or a human
        readable offset. Examples:

        - "next thursday at 3pm"
        - "tomorrow"
        - "3 days"
        - "2d"

        Times are in UTC.
        """

        timer = await self.create_timer(when.dt, 'reminder', ctx.author.id, ctx.channel.id, message)
        delta = time.human_timedelta(when.dt)
        await ctx.send(f"Alright {ctx.author.mention}, I'll remind you about {message} in {delta}.")

    async def on_reminder_timer_complete(self, timer):
        author_id, channel_id, message = timer.args

        channel = self.bot.get_channel(channel_id)
        if channel is None:
            # peculiar
            return

        await channel.send(f'<@{author_id}>, {timer.human_delta} you asked to be reminded of {message}.')

def setup(bot):
    bot.add_cog(Reminder(bot))
<commit_msg>Add ?remind alias for ?timer.<commit_after>from .utils import checks, db, time
from discord.ext import commands
import discord
import asyncio
import asyncpg
import datetime

class Reminders(db.Table):
    id = db.PrimaryKeyColumn()

    expires = db.Column(db.Datetime, index=True)
    created = db.Column(db.Datetime, default="now() at time zone 'utc'")
    event = db.Column(db.String)
    extra = db.Column(db.JSON, default="'{}'::jsonb")

class Timer:
    __slots__ = ('args', 'kwargs', 'event', 'id', 'created_at', 'expires')

    def __init__(self, *, record):
        self.id = record['id']

        extra = record['extra']
        self.args = extra.get('args', [])
        self.kwargs = extra.get('kwargs', {})
        self.event = record['event']
        self.created_at = record['created']
        self.expires = record['expires']

    @classmethod
    def temporary(cls, *, expires, created, event, args, kwargs):
        pseudo = {
            'id': None,
            'extra': { 'args': args, 'kwargs': kwargs },
            'event': event,
            'created': created,
            'expires': expires
        }
        return cls(record=pseudo)

    def __eq__(self, other):
        try:
            return self.id == other.id
        except AttributeError:
            return False

    def __hash__(self):
        return hash(self.id)

    @property
    def human_delta(self):
        return time.human_timedelta(self.created_at)

    def __repr__(self):
        return f'<Timer created={self.created_at} expires={self.expires} event={self.event}>'

class Reminder:
    """Reminders to do something."""

    def __init__(self, bot):
        self.bot = bot
        self._have_data = asyncio.Event(loop=bot.loop)
        self._current_timer = None
        self._task = bot.loop.create_task(self.dispatch_timers())

    def __unload(self):
        self._task.cancel()

    async def __error(self, ctx, error):
        if isinstance(error, commands.BadArgument):
            await ctx.send(error)

    async def get_active_timers(self, *, connection=None, days=7):
        query = "SELECT * FROM reminders WHERE expires < (CURRENT_DATE + $1::interval) ORDER BY expires;"
        con = connection or self.bot.pool

        records = await con.fetch(query, datetime.timedelta(days=days))
        return [Timer(record=a) for a in records]

    async def wait_for_active_timers(self, *, connection=None, days=7):
        async with db.MaybeAcquire(connection=connection, pool=self.bot.pool) as con:
            timers = await self.get_active_timers(connection=con, days=days)
            if len(timers):
                self._have_data.set()
                return timers

            self._have_data.clear()
            self._current_timer = None
            await self._have_data.wait()
            return await self.get_active_timers(connection=con, days=days)

    async def call_timer(self, timer):
        # delete the timer
        query = "DELETE FROM reminders WHERE id=$1;"
        await self.bot.pool.execute(query, timer.id)

        # dispatch the event
        event_name = f'{timer.event}_timer_complete'
        self.bot.dispatch(event_name, timer)

    async def dispatch_timers(self):
        try:
            while not self.bot.is_closed():
                # can only asyncio.sleep for up to ~48 days reliably
                # so we're gonna cap it off at 40 days
                # see: http://bugs.python.org/issue20493
                timers = await self.wait_for_active_timers(days=40)
                timer = self._current_timer = timers[0]
                now = datetime.datetime.utcnow()

                if timer.expires >= now:
                    to_sleep = (timer.expires - now).total_seconds()
                    await asyncio.sleep(to_sleep)

                await self.call_timer(timer)
        except asyncio.CancelledError:
            pass
        except (OSError, discord.ConnectionClosed, asyncpg.PostgresConnectionError):
            self._task.cancel()
            self._task = self.bot.loop.create_task(self.dispatch_timers())

    async def short_timer_optimisation(self, seconds, timer):
        await asyncio.sleep(seconds)
        event_name = f'{timer.event}_timer_complete'
        self.bot.dispatch(event_name, timer)

    async def create_timer(self, *args, **kwargs):
        """Creates a timer.

        Parameters
        -----------
        when: datetime.datetime
            When the timer should fire.
        event: str
            The name of the event to trigger.
            Will transform to 'on_{event}_timer_complete'.
        \*args
            Arguments to pass to the event
        \*\*kwargs
            Keyword arguments to pass to the event
        connection: asyncpg.Connection
            Special keyword-only argument to use a specific connection
            for the DB request.

        Note
        ------
        Arguments and keyword arguments must be JSON serialisable.

        Returns
        --------
        :class:`Timer`
        """
        when, event, *args = args

        try:
            connection = kwargs.pop('connection')
        except KeyError:
            connection = self.bot.pool

        now = datetime.datetime.utcnow()
        timer = Timer.temporary(event=event, args=args, kwargs=kwargs, expires=when, created=now)
        delta = (when - now).total_seconds()
        if delta <= 60:
            # a shortcut for small timers
            self.bot.loop.create_task(self.short_timer_optimisation(delta, timer))
            return timer

        query = """INSERT INTO reminders (event, extra, expires)
                   VALUES ($1, $2::jsonb, $3)
                   RETURNING id;
                """

        row = await connection.fetchrow(query, event, { 'args': args, 'kwargs': kwargs }, when)
        timer.id = row[0]

        self._have_data.set()

        # check if this timer is earlier than our currently run timer
        if self._current_timer and when < self._current_timer.expires:
            # cancel the task and re-run it
            self._task.cancel()
            self._task = self.bot.loop.create_task(self.dispatch_timers())

        return timer

    @commands.command(aliases=['timer', 'remind'])
    async def reminder(self, ctx, when: time.FutureTime, *, message: commands.clean_content = 'something'):
        """Reminds you of something after a certain amount of time.

        The time can be any direct date (e.g. YYYY-MM-DD) or a human
        readable offset. Examples:

        - "next thursday at 3pm"
        - "tomorrow"
        - "3 days"
        - "2d"

        Times are in UTC.
        """

        timer = await self.create_timer(when.dt, 'reminder', ctx.author.id, ctx.channel.id, message)
        delta = time.human_timedelta(when.dt)
        await ctx.send(f"Alright {ctx.author.mention}, I'll remind you about {message} in {delta}.")

    async def on_reminder_timer_complete(self, timer):
        author_id, channel_id, message = timer.args

        channel = self.bot.get_channel(channel_id)
        if channel is None:
            # peculiar
            return

        await channel.send(f'<@{author_id}>, {timer.human_delta} you asked to be reminded of {message}.')

def setup(bot):
    bot.add_cog(Reminder(bot))
<|endoftext|>