<commit_before>from io import BytesIO

from discord import File, Member, User
from discord.ext import commands
from discord.ext.commands import Cog

from bot import BeattieBot
from context import BContext

class Default(Cog):
    """Default useful commands."""

    @commands.command()
    async def avatar(
        self, ctx: BContext, *, user: User = commands.param(default=commands.Author)
    ) -> None:
        img = BytesIO()
        avatar = user.display_avatar.with_format(
            "gif" if user.display_avatar.is_animated() else "png"
        )
        await avatar.save(img)
        filename = str(avatar).rpartition("/")[2].partition("?")[0]
        await ctx.send(file=File(img, filename))

    @avatar.error  # type: ignore
    async def avatar_error(self, ctx: BContext, exc: Exception) -> None:
        if isinstance(exc, commands.BadArgument):
            await ctx.send("User not found.")
        else:
            await ctx.bot.handle_error(ctx, exc)

    @commands.command()
    async def latency(self, ctx: BContext) -> None:
        """Get the latency to the websocket."""
        await ctx.send(f"WS latency: **{ctx.bot.latency*1000:.0f}ms**")

    @commands.command()
    async def ping(self, ctx: BContext) -> None:
        """Get how fast the bot reacts to a command message"""
        msg = await ctx.send("...")
        delta = msg.created_at - ctx.message.created_at
        await msg.edit(content=f":ping_pong: **{delta.total_seconds()*1000:.0f}ms**")
        msg = await ctx.channel.fetch_message(msg.id)
        edited_at = msg.edited_at
        assert edited_at is not None
        delta2 = edited_at - ctx.message.created_at
        await msg.edit(
            content=f"{msg.content}\n**{delta2.total_seconds()*1000:.0f}ms**"
        )

    @commands.command()
    async def source(self, ctx: BContext) -> None:
        """Get the source for the bot."""
        await ctx.send("https://github.com/BeatButton/beattie")

    @commands.command()
    async def invite(self, ctx: BContext) -> None:
        """Get the invite for the bot."""
        url = "<https://discordapp.com/oauth2/authorize?client_id={}&scope=bot>"
        await ctx.send(url.format(ctx.me.id))

async def setup(bot: BeattieBot) -> None:
    await bot.add_cog(Default())
<commit_msg>fix avatar default argument<commit_after>from io import BytesIO

from discord import File, Member, User
from discord.ext import commands
from discord.ext.commands import Cog

from bot import BeattieBot
from context import BContext

class Default(Cog):
    """Default useful commands."""

    @commands.command()
    async def avatar(self, ctx: BContext, *, user: User = commands.Author) -> None:
        img = BytesIO()
        avatar = user.display_avatar.with_format(
            "gif" if user.display_avatar.is_animated() else "png"
        )
        await avatar.save(img)
        filename = str(avatar).rpartition("/")[2].partition("?")[0]
        await ctx.send(file=File(img, filename))

    @avatar.error  # type: ignore
    async def avatar_error(self, ctx: BContext, exc: Exception) -> None:
        if isinstance(exc, commands.BadArgument):
            await ctx.send("User not found.")
        else:
            await ctx.bot.handle_error(ctx, exc)

    @commands.command()
    async def latency(self, ctx: BContext) -> None:
        """Get the latency to the websocket."""
        await ctx.send(f"WS latency: **{ctx.bot.latency*1000:.0f}ms**")

    @commands.command()
    async def ping(self, ctx: BContext) -> None:
        """Get how fast the bot reacts to a command message"""
        msg = await ctx.send("...")
        delta = msg.created_at - ctx.message.created_at
        await msg.edit(content=f":ping_pong: **{delta.total_seconds()*1000:.0f}ms**")
        msg = await ctx.channel.fetch_message(msg.id)
        edited_at = msg.edited_at
        assert edited_at is not None
        delta2 = edited_at - ctx.message.created_at
        await msg.edit(
            content=f"{msg.content}\n**{delta2.total_seconds()*1000:.0f}ms**"
        )

    @commands.command()
    async def source(self, ctx: BContext) -> None:
        """Get the source for the bot."""
        await ctx.send("https://github.com/BeatButton/beattie")

    @commands.command()
    async def invite(self, ctx: BContext) -> None:
        """Get the invite for the bot."""
        url = "<https://discordapp.com/oauth2/authorize?client_id={}&scope=bot>"
        await ctx.send(url.format(ctx.me.id))

async def setup(bot: BeattieBot) -> None:
    await bot.add_cog(Default())
<|endoftext|>


<commit_before>"""
This cog is to be used primarily for small janitorial tasks
(removing clover once member is hit, pruning clovers)
"""
import discord
from discord import AuditLogAction
from discord.ext import commands
from .utils import checks
from .utils.enums import Change
from datetime import datetime, timedelta
import asyncio

class Janitor():
    """
    The main class wrapepr
    """
    def __init__(self, bot):
        super().__init__()
        self.bot = bot
        self.owner = None
        try:
            self.bg_task = self.bot.loop.create_task(self.daily_prune())
            self.owner_task = self.bot.loop.create_task(self.setup_owner_dm())
        except Exception as e:
            self.bot.logger.warning(f'Error starting task prune_clovers: {e}')

    async def setup_owner_dm(self):
        await self.bot.wait_until_ready()
        self.bot.logger.info(
            f'Setting up owner channel {self.bot.bot_owner_id}')
        try:
            self.server_logs = self.bot.get_channel(378684962934751239)
            self.owner = await self.bot.get_user_info(self.bot.bot_owner_id)
        except Exception as e:
            self.bot.logger.warning(f'Error getting owner: {e}')
        self.bot.logger.info(f'User retrieved: {self.owner.name}')
        try:
            try:
                await self.owner.create_dm()
            except Exception as e:
                self.bot.logger.warning(f'Error creating dm channel: {e}')
            await self.owner.dm_channel.send('Bot started successfully')
            #await self.server_logs.send('Bot started successfully')
        except Exception as e:
            self.bot.logger.warning(f'Error getting dm channel: {e}')

    def remove_clover(self, member) -> list:
        member_roles = member.roles
        for index, role in enumerate(member_roles):
            if role.name.lower() == 'clover' or role.name.lower() == 'dedicated':
                del member_roles[index]
        return member_roles

    def remove_key(self, member) -> list:
        member_roles = member.roles.copy
        for index, role in enumerate(member_roles):
            if role.name.lower() == '🔑':
                key_index = index
        del member_roles[key_index]
        return member_roles

    def remove_access(self, member) -> list:
        member_roles = member.roles.copy()
        role_list = []
        for index, role in enumerate(member_roles):
            if role.name.lower() in ['muse', 'dev','lewd', 'swole', 'artsy', 'shokugeki', 'degenerate', 'simulcast', 'legacy', 'meta', 'stylish']:
                role_list.append(role)
        for role in role_list:
            member_roles.remove(role)
        return member_roles

    async def on_message(self, message):
        if message.author.bot:
            return
        if message.content.startswith('.iam'):
            return
        has_clover = False
        has_key = False
        has_member = False
        has_dedicated = False
        has_permrole = False
        member_roles = message.author.roles
        for index, role in enumerate(member_roles):
            if role.name.lower() == 'clover':
                has_clover = True
            elif role.name.lower() == '🔑':
                has_key = True
            elif role.name.lower() == 'dedicated':
                has_dedicated = True
            elif role.name.lower() in ['muse', 'dev','lewd', 'swole', 'artsy', 'shokugeki', 'degenerate', 'simulcast', 'legacy', 'meta', 'stylish']:
                has_permrole = True
            elif role.name.lower() == 'member':
                has_member = True
        if has_clover and has_member:
            member_roles = self.remove_clover(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User upgraded from clover to member")
                await message.add_reaction('🎊')
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' was just promoted to member!')
                try:
                    await self.bot.postgres_controller.insert_rolechange(
                        message.guild.id, message.author.id, Change.PROMOTION
                    )
                except Exception as e:
                    self.bot.logger.warning(
                        f'Issue logging action to db: {e})')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users roles: {e}')
        if has_key and has_member:
            member_roles = self.remove_key(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User upgraded from Key to member")
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' was just promoted to member!')
                try:
                    await self.bot.postgres_controller.insert_rolechange(
                        message.guild.id, message.author.id, Change.PROMOTION
                    )
                except Exception as e:
                    self.bot.logger.warning(
                        f'Issue logging action to db: {e})')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users roles: {e}')
        if has_dedicated and has_permrole:
            member_roles = self.remove_access(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User had perm roles")
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' had access roles removed')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users access roles: {e}')

    @commands.command(hidden=True)
    @checks.has_permissions(manage_roles=True)
    async def aggroprune(self,ctx):
        self.bot.logger.info(f'Prune requested by: {ctx.message.author}')
        await self.prune_nonclovers()

    @commands.command(hidden=True)
    @checks.has_permissions(manage_roles=True)
    async def prune(self, ctx):
        self.bot.logger.info(f'Prune requested by: {ctx.message.author}')
        await self.prune_clovers()

    @prune.error
    async def prune_error(self, ctx, error):
        if isinstance(error, commands.errors.CheckFailure):
            self.bot.logger.warning(f'{ctx.message.author} '
                                    'tried to run prune w/o permissions')

    @aggroprune.error
    async def aggroprune_error(self, ctx, error):
        if isinstance(error, commands.errors.CheckFailure):
            self.bot.logger.warning(f'{ctx.message.author} '
                                    'tried to run prune w/o permissions')

    async def daily_prune(self):
        self.bot.logger.info("Starting prune task, first prune in 24 hours")
        while not self.bot.is_closed():
            await asyncio.sleep(86400)
            await self.prune_clovers()

    async def prune_clovers(self):
        try:
            self.bot.logger.info('Starting prune task now')
        except Exception as e:
            self.bot.logger.info('tf')
        clovers = []
        clover_role = None
        mod_log = self.bot.get_channel(self.bot.mod_log)
        a_irl = self.bot.get_guild(self.bot.guild_id)
        for role in a_irl.roles:
            if role.name.lower() == 'clover':
                clover_role = role
        if not clover_role:
            self.bot.logger.warning(
                'Something went really wrong, '
                'I couldn\'t find the clover role')
            return
        clovers = clover_role.members
        try:
            members_prunable = await self.bot.postgres_controller.get_all_prunable()
        except Exception as e:
            self.bot.logger.warning(f'{e}')
        prune_info = {'pruned': False, 'amount': 0}
        self.bot.logger.info(f'{members_prunable}')
        for member in clovers:
            if member.id in members_prunable:
                try:
                    new_roles = self.remove_clover(member)
                    await member.edit(
                        roles=new_roles,
                        reason="Pruned due to inactivity"
                    )
                    prune_info['pruned'] = True
                    prune_info['amount'] += 1
                except Exception as e:
                    self.bot.logger.warning(
                        f'Error pruning clovers: {e}'
                    )
        self.bot.logger.info(f'Prune info: {prune_info}')
        try:
            await self.bot.postgres_controller.reset_message_deleted()
        except Exception as e:
            self.bot.logger.warning(f'Issue resetting spam db: {e}')
        if prune_info['pruned']:
            try:
                await mod_log.send(
                    f'Pruned {prune_info["amount"]} clovers 🍀🔫')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error posting prune info to mod_log: {e}')

    async def prune_nonclovers(self):
        try:
            self.bot.logger.info('Starting aggroprune task now')
        except Exception as e:
            self.bot.logger.info('tf')
        clovers = []
        clover_role = None
        mod_log = self.bot.get_channel(self.bot.mod_log)
        a_irl = self.bot.get_guild(self.bot.guild_id)
        for role in a_irl.roles:
            if role.name.lower() == 'clover':
                clover_role = role
        if not clover_role:
            self.bot.logger.warning(
                'Something went really wrong, '
                'I couldn\'t find the clover role')
            return
        clovers = clover_role.members
        self.bot.logger.info(f'Can i get uhhh')
        try:
            members_safe = await self.bot.postgres_controller.get_all_clovers()
        except Exception as e:
            self.bot.logger.warning(f'shiiiiittt:{e}')
        self.bot.logger.info(f'uhhh {members_safe}')
        prune_info = {'pruned': False, 'amount': 0}
        for member in clovers:
            if member.id not in members_safe:
                try:
                    new_roles = self.remove_clover(member)
                    await member.edit(
                        roles=new_roles,
                        reason="Pruned due to inactivity"
                    )
                    prune_info['pruned'] = True
                    prune_info['amount'] += 1
                except Exception as e:
                    self.bot.logger.warning(
                        f'Error pruning clovers: {e}'
                    )
        self.bot.logger.info(f'Prune info: {prune_info}')
        try:
            await self.bot.postgres_controller.reset_message_deleted()
            await self.bot.postgres_controller.reset_user_reactions()
        except Exception as e:
            self.bot.logger.warning(f'Issue resetting spam db: {e}')
        if prune_info['pruned']:
            try:
                await mod_log.send(
                    f'Pruned {prune_info["amount"]} clovers 🍀🔫')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error posting prune info to mod_log: {e}')

    @commands.command()
    async def month_end(self, ctx):
        roles_to_wipe = ['member', 'active', 'regular', 'contributor', 'addicted', 'insomniac', 'no-lifer']
        local_embed = discord.Embed(
            title=f'Resetting Roles...',
            description=f'Please be patient, this might take a while...'
        )
        value_string =f'Clearing :key: role ... \n'\
                      f'Resetting *Member* role ...\n'\
                      f'Resetting *Active* role ...\n'\
                      f'Resetting *Regular* role ...\n'\
                      f'Resetting *Contributor* role ...\n'\
                      f'Resetting *Addicted* role ...\n'\
                      f'Resetting *Insomniac* role ...\n'\
                      f'Resetting *No-Lifer* role ...\n'
        local_embed.add_field(name="Progress:",
                value=value_string)
        values = value_string.split('\n')
        for index, val in enumerate(values):
            values[index] = val +'\n'
        message = await ctx.send(embed=local_embed)
        key_role = None
        for role in ctx.channel.guild.roles:
            if role.name.lower() == '🔑':
                key_role = role
        try:
            await self.rem_all_members(ctx, key_role)
            values[0] = f'Clearing :key: role :white_check_mark:'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0 , name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)
        except:
            values[0] = f'Clearing :key: role :x:'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0, name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)
        for counter, role in enumerate(roles_to_wipe):
            members = await self.get_all_members(ctx, role)
            if not members:
                values[counter+1] = f'Resetting **{role.title()}** role :x:'
                l_embed = message.embeds[0]
                l_embed.set_field_at(0, name="Progress:", value="".join(values))
                await message.edit(embed = l_embed)
                self.bot.logger.warning(f'ayy yo i didn"t find this one {role}')
            for member in members:
                temp_roles = self.rem_role(member, role)
                temp_roles = self.add_role(temp_roles, key_role)
                await member.edit(roles=temp_roles)
            values[counter+1] = f'Resetting **{role.title()}** role :white_check_mark:'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0, name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)

    async def rem_all_members(self, ctx, rem_role):
        for member in rem_role.members:
            roles_list = member.roles.copy()
            for counter, role in enumerate(roles_list):
                if role == rem_role:
                    roles_list = roles_list.pop(counter)
                    await member.edit(roles=roles_list)
                    continue

    async def get_all_members(self, ctx, role_name):
        found_role = None
        for role in ctx.channel.guild.roles:
            if role.name.lower() == role_name:
                found_role = role
        if not found_role:
            return None
        return found_role.members

    def add_role(self, role_list, role):
        role_list.append(role)
        return role_list

    def rem_role(self, member, role) -> list:
        member_roles = member.roles.copy()
        for index, role in enumerate(member_roles):
            if role.name.lower() == role:
                key_index = index
        del member_roles[key_index]
        return member_roles
<commit_msg>newlines<commit_after>"""
This cog is to be used primarily for small janitorial tasks
(removing clover once member is hit, pruning clovers)
"""
import discord
from discord import AuditLogAction
from discord.ext import commands
from .utils import checks
from .utils.enums import Change
from datetime import datetime, timedelta
import asyncio

class Janitor():
    """
    The main class wrapepr
    """
    def __init__(self, bot):
        super().__init__()
        self.bot = bot
        self.owner = None
        try:
            self.bg_task = self.bot.loop.create_task(self.daily_prune())
            self.owner_task = self.bot.loop.create_task(self.setup_owner_dm())
        except Exception as e:
            self.bot.logger.warning(f'Error starting task prune_clovers: {e}')

    async def setup_owner_dm(self):
        await self.bot.wait_until_ready()
        self.bot.logger.info(
            f'Setting up owner channel {self.bot.bot_owner_id}')
        try:
            self.server_logs = self.bot.get_channel(378684962934751239)
            self.owner = await self.bot.get_user_info(self.bot.bot_owner_id)
        except Exception as e:
            self.bot.logger.warning(f'Error getting owner: {e}')
        self.bot.logger.info(f'User retrieved: {self.owner.name}')
        try:
            try:
                await self.owner.create_dm()
            except Exception as e:
                self.bot.logger.warning(f'Error creating dm channel: {e}')
            await self.owner.dm_channel.send('Bot started successfully')
            #await self.server_logs.send('Bot started successfully')
        except Exception as e:
            self.bot.logger.warning(f'Error getting dm channel: {e}')

    def remove_clover(self, member) -> list:
        member_roles = member.roles
        for index, role in enumerate(member_roles):
            if role.name.lower() == 'clover' or role.name.lower() == 'dedicated':
                del member_roles[index]
        return member_roles

    def remove_key(self, member) -> list:
        member_roles = member.roles.copy
        for index, role in enumerate(member_roles):
            if role.name.lower() == '🔑':
                key_index = index
        del member_roles[key_index]
        return member_roles

    def remove_access(self, member) -> list:
        member_roles = member.roles.copy()
        role_list = []
        for index, role in enumerate(member_roles):
            if role.name.lower() in ['muse', 'dev','lewd', 'swole', 'artsy', 'shokugeki', 'degenerate', 'simulcast', 'legacy', 'meta', 'stylish']:
                role_list.append(role)
        for role in role_list:
            member_roles.remove(role)
        return member_roles

    async def on_message(self, message):
        if message.author.bot:
            return
        if message.content.startswith('.iam'):
            return
        has_clover = False
        has_key = False
        has_member = False
        has_dedicated = False
        has_permrole = False
        member_roles = message.author.roles
        for index, role in enumerate(member_roles):
            if role.name.lower() == 'clover':
                has_clover = True
            elif role.name.lower() == '🔑':
                has_key = True
            elif role.name.lower() == 'dedicated':
                has_dedicated = True
            elif role.name.lower() in ['muse', 'dev','lewd', 'swole', 'artsy', 'shokugeki', 'degenerate', 'simulcast', 'legacy', 'meta', 'stylish']:
                has_permrole = True
            elif role.name.lower() == 'member':
                has_member = True
        if has_clover and has_member:
            member_roles = self.remove_clover(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User upgraded from clover to member")
                await message.add_reaction('🎊')
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' was just promoted to member!')
                try:
                    await self.bot.postgres_controller.insert_rolechange(
                        message.guild.id, message.author.id, Change.PROMOTION
                    )
                except Exception as e:
                    self.bot.logger.warning(
                        f'Issue logging action to db: {e})')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users roles: {e}')
        if has_key and has_member:
            member_roles = self.remove_key(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User upgraded from Key to member")
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' was just promoted to member!')
                try:
                    await self.bot.postgres_controller.insert_rolechange(
                        message.guild.id, message.author.id, Change.PROMOTION
                    )
                except Exception as e:
                    self.bot.logger.warning(
                        f'Issue logging action to db: {e})')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users roles: {e}')
        if has_dedicated and has_permrole:
            member_roles = self.remove_access(message.author)
            try:
                await message.author.edit(
                    roles=member_roles,
                    reason="User had perm roles")
                self.bot.logger.info(
                    f'{message.author.display_name}'
                    ' had access roles removed')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error updating users access roles: {e}')

    @commands.command(hidden=True)
    @checks.has_permissions(manage_roles=True)
    async def aggroprune(self,ctx):
        self.bot.logger.info(f'Prune requested by: {ctx.message.author}')
        await self.prune_nonclovers()

    @commands.command(hidden=True)
    @checks.has_permissions(manage_roles=True)
    async def prune(self, ctx):
        self.bot.logger.info(f'Prune requested by: {ctx.message.author}')
        await self.prune_clovers()

    @prune.error
    async def prune_error(self, ctx, error):
        if isinstance(error, commands.errors.CheckFailure):
            self.bot.logger.warning(f'{ctx.message.author} '
                                    'tried to run prune w/o permissions')

    @aggroprune.error
    async def aggroprune_error(self, ctx, error):
        if isinstance(error, commands.errors.CheckFailure):
            self.bot.logger.warning(f'{ctx.message.author} '
                                    'tried to run prune w/o permissions')

    async def daily_prune(self):
        self.bot.logger.info("Starting prune task, first prune in 24 hours")
        while not self.bot.is_closed():
            await asyncio.sleep(86400)
            await self.prune_clovers()

    async def prune_clovers(self):
        try:
            self.bot.logger.info('Starting prune task now')
        except Exception as e:
            self.bot.logger.info('tf')
        clovers = []
        clover_role = None
        mod_log = self.bot.get_channel(self.bot.mod_log)
        a_irl = self.bot.get_guild(self.bot.guild_id)
        for role in a_irl.roles:
            if role.name.lower() == 'clover':
                clover_role = role
        if not clover_role:
            self.bot.logger.warning(
                'Something went really wrong, '
                'I couldn\'t find the clover role')
            return
        clovers = clover_role.members
        try:
            members_prunable = await self.bot.postgres_controller.get_all_prunable()
        except Exception as e:
            self.bot.logger.warning(f'{e}')
        prune_info = {'pruned': False, 'amount': 0}
        self.bot.logger.info(f'{members_prunable}')
        for member in clovers:
            if member.id in members_prunable:
                try:
                    new_roles = self.remove_clover(member)
                    await member.edit(
                        roles=new_roles,
                        reason="Pruned due to inactivity"
                    )
                    prune_info['pruned'] = True
                    prune_info['amount'] += 1
                except Exception as e:
                    self.bot.logger.warning(
                        f'Error pruning clovers: {e}'
                    )
        self.bot.logger.info(f'Prune info: {prune_info}')
        try:
            await self.bot.postgres_controller.reset_message_deleted()
        except Exception as e:
            self.bot.logger.warning(f'Issue resetting spam db: {e}')
        if prune_info['pruned']:
            try:
                await mod_log.send(
                    f'Pruned {prune_info["amount"]} clovers 🍀🔫')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error posting prune info to mod_log: {e}')

    async def prune_nonclovers(self):
        try:
            self.bot.logger.info('Starting aggroprune task now')
        except Exception as e:
            self.bot.logger.info('tf')
        clovers = []
        clover_role = None
        mod_log = self.bot.get_channel(self.bot.mod_log)
        a_irl = self.bot.get_guild(self.bot.guild_id)
        for role in a_irl.roles:
            if role.name.lower() == 'clover':
                clover_role = role
        if not clover_role:
            self.bot.logger.warning(
                'Something went really wrong, '
                'I couldn\'t find the clover role')
            return
        clovers = clover_role.members
        self.bot.logger.info(f'Can i get uhhh')
        try:
            members_safe = await self.bot.postgres_controller.get_all_clovers()
        except Exception as e:
            self.bot.logger.warning(f'shiiiiittt:{e}')
        self.bot.logger.info(f'uhhh {members_safe}')
        prune_info = {'pruned': False, 'amount': 0}
        for member in clovers:
            if member.id not in members_safe:
                try:
                    new_roles = self.remove_clover(member)
                    await member.edit(
                        roles=new_roles,
                        reason="Pruned due to inactivity"
                    )
                    prune_info['pruned'] = True
                    prune_info['amount'] += 1
                except Exception as e:
                    self.bot.logger.warning(
                        f'Error pruning clovers: {e}'
                    )
        self.bot.logger.info(f'Prune info: {prune_info}')
        try:
            await self.bot.postgres_controller.reset_message_deleted()
            await self.bot.postgres_controller.reset_user_reactions()
        except Exception as e:
            self.bot.logger.warning(f'Issue resetting spam db: {e}')
        if prune_info['pruned']:
            try:
                await mod_log.send(
                    f'Pruned {prune_info["amount"]} clovers 🍀🔫')
            except Exception as e:
                self.bot.logger.warning(
                    f'Error posting prune info to mod_log: {e}')

    @commands.command()
    async def month_end(self, ctx):
        roles_to_wipe = ['member', 'active', 'regular', 'contributor', 'addicted', 'insomniac', 'no-lifer']
        local_embed = discord.Embed(
            title=f'Resetting Roles...',
            description=f'Please be patient, this might take a while...'
        )
        value_string =f'Clearing :key: role ... \n'\
                      f'Resetting *Member* role ...\n'\
                      f'Resetting *Active* role ...\n'\
                      f'Resetting *Regular* role ...\n'\
                      f'Resetting *Contributor* role ...\n'\
                      f'Resetting *Addicted* role ...\n'\
                      f'Resetting *Insomniac* role ...\n'\
                      f'Resetting *No-Lifer* role ...\n'
        local_embed.add_field(name="Progress:",
                value=value_string)
        values = value_string.split('\n')
        for index, val in enumerate(values):
            values[index] = val +'\n'
        message = await ctx.send(embed=local_embed)
        key_role = None
        for role in ctx.channel.guild.roles:
            if role.name.lower() == '🔑':
                key_role = role
        try:
            await self.rem_all_members(ctx, key_role)
            values[0] = f'Clearing :key: role :white_check_mark: \n'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0 , name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)
        except:
            values[0] = f'Clearing :key: role :x:'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0, name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)
        for counter, role in enumerate(roles_to_wipe):
            members = await self.get_all_members(ctx, role)
            if not members:
                values[counter+1] = f'Resetting **{role.title()}** role :x:\n'
                l_embed = message.embeds[0]
                l_embed.set_field_at(0, name="Progress:", value="".join(values))
                await message.edit(embed = l_embed)
                self.bot.logger.warning(f'ayy yo i didn"t find this one {role}')
            for member in members:
                temp_roles = self.rem_role(member, role)
                temp_roles = self.add_role(temp_roles, key_role)
                await member.edit(roles=temp_roles)
            values[counter+1] = f'Resetting **{role.title()}** role :white_check_mark:\n'
            l_embed = message.embeds[0]
            l_embed.set_field_at(0, name="Progress:", value="".join(values))
            await message.edit(embed = l_embed)

    async def rem_all_members(self, ctx, rem_role):
        for member in rem_role.members:
            roles_list = member.roles.copy()
            for counter, role in enumerate(roles_list):
                if role == rem_role:
                    roles_list = roles_list.pop(counter)
                    await member.edit(roles=roles_list)
                    continue

    async def get_all_members(self, ctx, role_name):
        found_role = None
        for role in ctx.channel.guild.roles:
            if role.name.lower() == role_name:
                found_role = role
        if not found_role:
            return None
        return found_role.members

    def add_role(self, role_list, role):
        role_list.append(role)
        return role_list

    def rem_role(self, member, role) -> list:
        member_roles = member.roles.copy()
        for index, role in enumerate(member_roles):
            if role.name.lower() == role:
                key_index = index
        del member_roles[key_index]
        return member_roles
<|endoftext|>


<commit_before>from django import forms
from django.core.urlresolvers import reverse_lazy

from resumable.fields import ResumableFileField

from sections.models import SubSubSection

class ResumableForm(forms.Form):
    title = forms.CharField()
    file_name = ResumableFileField(upload_url=reverse_lazy('upload'), chunks_dir="chuncks")
    subsubsection = forms.ModelChoiceField(queryset=SubSubSection.objects.filter(subsection__section="1"), required=False)
<commit_msg>[enh] order subsubsection by subsection.title in upload video<commit_after>from django import forms
from django.core.urlresolvers import reverse_lazy

from resumable.fields import ResumableFileField

from sections.models import SubSubSection

class ResumableForm(forms.Form):
    title = forms.CharField()
    file_name = ResumableFileField(upload_url=reverse_lazy('upload'), chunks_dir="chuncks")
    subsubsection = forms.ModelChoiceField(queryset=SubSubSection.objects.filter(subsection__section="1").order_by("subsection__title"), required=False)
<|endoftext|>


<commit_before>import discord
from discord.ext import commands
from __main__ import settings, botdata, httpgetter
from cogs.utils.helpers import *
from cogs.utils.clip import *
from cogs.utils import checks
from .mangocog import *
import os
import re

async def pokeapi_query(url, fullurl=False):
	if not fullurl:
		url = f"http://pokeapi.co/api/v2{url}"
	return await httpgetter.get(url, cache=True, errors={
		404: "Pokemon not found",
		"default": "pokeapi said we did things wrong 😢. status code: {}"
	})

def poke_color(color):
	return {
		"black": discord.Color(0x000000),
		"blue": discord.Color.blue(),
		"brown": discord.Color(0xD2691E),
		"gray": discord.Color(0xA9A9A9),
		"green": discord.Color.green(),
		"pink": discord.Color(0xFF69B4),
		"purple": discord.Color.purple(),
		"red": discord.Color.red(),
		"white": discord.Color(0xFFFFFF),
		"yellow": discord.Color(0xFFFF00)
	}[color]

def localize(list_data, end_key):
	for item in list_data:
		if item["language"]["name"] == "en":
			return item[end_key]
	return None

class Pokemon(MangoCog):
	"""Pokemon related commands

	A few commands using the [pokeapi](https://pokeapi.co/) which provides information about pokemon

	Note that the above API does not include information about pokemon Sun and Moon, so I can't look up those pokemon for you"""
	def __init__(self, bot):
		MangoCog.__init__(self, bot)

	def poke_type(self, type_name):
		if type_name == "shadow" or type_name == "unknown":
			return type_name
		else:
			return self.get_emoji(f"poke_{type_name}")

	async def get_pokemon_data(self, pokemon):
		# Sanitize input first
		pokemon = pokemon.lower()
		replacements = { " ": "-", "♂": "-m", "♀": "-f" }
		for key in replacements:
			pokemon = pokemon.replace(key, replacements[key])
		pokemon = re.sub(r'[^a-z0-9\-]', '', pokemon)

		words = pokemon.split("-")
		if "mega" in words:
			words.remove("mega")
			words.insert(1, "mega")
			pokemon = "-".join(words)

		try:
			form_data = await pokeapi_query(f"/pokemon-form/{pokemon}/")
			data = await pokeapi_query(form_data["pokemon"]["url"], True)
		except Http404Error as e:
			form_data = None
			data = await pokeapi_query(f"/pokemon/{pokemon}/")
		species_data = await pokeapi_query(data["species"]["url"], True)

		data["localized_name"] = localize(species_data["names"], "name")
		if data["id"] >= 1000:
			data["id"] = species_data["id"]
		data["wiki_url"] = f"http://www.serebii.net/pokedex-sm/{data['id']:03d}.shtml"

		if form_data:
			data["sprites"] = form_data["sprites"]
			name = localize(form_data["names"], "name")
			if name:
				data["localized_name"] = name
			if form_data.get("is_mega"):
				data["wiki_url"] += "#mega"

		return data, species_data

	# returns True on success, False on failure
	async def play_pokecry(self, ctx, poke_id, namestring, old=False):
		is_mega = "mega" in str(namestring).lower()

		name = str(poke_id)
		if old:
			name = f"old/{name}"
		if is_mega:
			name = f"{name}-mega"

		try:
			clip = await self.get_clip(f"url:http://dillerm.io/pokemon/cries/{name}.ogg", ctx)
			clip.volume = 0.1
			await self.play_clip(clip, ctx)
		except Http404Error:
			return False

		return True

	@commands.command(aliases=["pokemon"])
	async def pokedex(self, ctx, *, pokemon):
		"""Looks up information about the indicated pokemon

		Pokemon should be specified using either their name or id number

		Clicking on the pokemon's name will bring you to their wiki page

		Example:
		`{cmdpfx}pokedex charizard`"""
		with ctx.channel.typing():
			data, species_data = await self.get_pokemon_data(pokemon)

		types = []
		for t in sorted(data["types"], key=lambda t: t["slot"]):
			types.append(self.poke_type(t["type"]["name"]))

		flavor_text = localize(species_data["flavor_text_entries"], "flavor_text")
		flavor_text = flavor_text.replace("\n", " ")

		embed = discord.Embed(description=flavor_text, color=poke_color(species_data["color"]["name"]))
		embed.title = data["localized_name"] + f" #{data['id']}"
		embed.url = data["wiki_url"]

		embed.set_thumbnail(url=data["sprites"]["front_default"])

		embed.add_field(name=f"Type{'s' if len(types) > 1 else ''}", value=f"{''.join(types)}")
		if species_data.get("habitat"):
			embed.add_field(name="Habitat", value=f"{species_data['habitat']['name']}")
		embed.add_field(name="Weight", value=f"{data['weight'] / 10} kg")
		embed.add_field(name="Height", value=f"{data['height'] / 10} m")

		# fails silently if there's no cry for this pokemon
		await self.play_pokecry(ctx, data["id"], pokemon)

		await ctx.send(embed=embed)

	@commands.command()
	async def shiny(self, ctx, *, pokemon):
		"""Gets the shiny version of this pokemon

		pokemon should be specified using either their name or id number

		Example:
		`{cmdpfx}shiny charizard`"""

		# Sanitize input first
		with ctx.channel.typing():
			data, species_data = await self.get_pokemon_data(pokemon)

		if not data["sprites"].get("front_shiny"):
			await ctx.send("This pokemon doesn't have a shiny version")

		embed = discord.Embed(color=poke_color(species_data["color"]["name"]))
		embed.set_image(url=data["sprites"].get("front_shiny"))

		# fails silently if there's no cry for this pokemon
		await self.play_pokecry(ctx, data["id"], pokemon)

		await ctx.send(embed=embed)

	@commands.command(aliases=["cry"])
	async def pokecry(self, ctx, *, pokemon):
		"""Plays the pokemon's sound effect

		Audio files for these pokemon cries were gotten from [Veekun](https://veekun.com/dex/downloads). Veekun does not have the cries for Generation VII yet, so I won't be able to play those.

		Most pokemon have a new cry and an old cry. To get the old cry instead of the new one, add 'old' to the end of your command (see example below.)

		**Example:**
		`{cmdpfx}pokecry pikachu`
		`{cmdpfx}pokecry bulbasaur old`"""
		words = pokemon.split(" ")
		old = "old" in words
		if old:
			words.remove("old")
		pokemon = " ".join(words)

		data, species_data = await self.get_pokemon_data(pokemon)

		if data["id"] > 721:
			raise UserError("Sorry, I don't have the cries for pokemon in Generation VII yet")

		success = await self.play_pokecry(ctx, data["id"], pokemon, old=old)

		if not success:
			raise UserError(f"Couldn't find the cry for {data['localized_name']}")

def setup(bot):
	bot.add_cog(Pokemon(bot))<commit_msg>fixed pokemon commands to work when mango is not in a voice channel<commit_after>import discord
from discord.ext import commands
from __main__ import settings, botdata, httpgetter
from cogs.utils.helpers import *
from cogs.utils.clip import *
from cogs.utils import checks
from .mangocog import *
from cogs.audio import AudioPlayerNotFoundError
import os
import re

async def pokeapi_query(url, fullurl=False):
	if not fullurl:
		url = f"http://pokeapi.co/api/v2{url}"
	return await httpgetter.get(url, cache=True, errors={
		404: "Pokemon not found",
		"default": "pokeapi said we did things wrong 😢. status code: {}"
	})

def poke_color(color):
	return {
		"black": discord.Color(0x000000),
		"blue": discord.Color.blue(),
		"brown": discord.Color(0xD2691E),
		"gray": discord.Color(0xA9A9A9),
		"green": discord.Color.green(),
		"pink": discord.Color(0xFF69B4),
		"purple": discord.Color.purple(),
		"red": discord.Color.red(),
		"white": discord.Color(0xFFFFFF),
		"yellow": discord.Color(0xFFFF00)
	}[color]

def localize(list_data, end_key):
	for item in list_data:
		if item["language"]["name"] == "en":
			return item[end_key]
	return None

class Pokemon(MangoCog):
	"""Pokemon related commands

	A few commands using the [pokeapi](https://pokeapi.co/) which provides information about pokemon

	Note that the above API does not include information about pokemon Sun and Moon, so I can't look up those pokemon for you"""
	def __init__(self, bot):
		MangoCog.__init__(self, bot)

	def poke_type(self, type_name):
		if type_name == "shadow" or type_name == "unknown":
			return type_name
		else:
			return self.get_emoji(f"poke_{type_name}")

	async def get_pokemon_data(self, pokemon):
		# Sanitize input first
		pokemon = pokemon.lower()
		replacements = { " ": "-", "♂": "-m", "♀": "-f" }
		for key in replacements:
			pokemon = pokemon.replace(key, replacements[key])
		pokemon = re.sub(r'[^a-z0-9\-]', '', pokemon)

		words = pokemon.split("-")
		if "mega" in words:
			words.remove("mega")
			words.insert(1, "mega")
			pokemon = "-".join(words)

		try:
			form_data = await pokeapi_query(f"/pokemon-form/{pokemon}/")
			data = await pokeapi_query(form_data["pokemon"]["url"], True)
		except Http404Error as e:
			form_data = None
			data = await pokeapi_query(f"/pokemon/{pokemon}/")
		species_data = await pokeapi_query(data["species"]["url"], True)

		data["localized_name"] = localize(species_data["names"], "name")
		if data["id"] >= 1000:
			data["id"] = species_data["id"]
		data["wiki_url"] = f"http://www.serebii.net/pokedex-sm/{data['id']:03d}.shtml"

		if form_data:
			data["sprites"] = form_data["sprites"]
			name = localize(form_data["names"], "name")
			if name:
				data["localized_name"] = name
			if form_data.get("is_mega"):
				data["wiki_url"] += "#mega"

		return data, species_data

	# returns True on success, False on failure
	async def play_pokecry(self, ctx, poke_id, namestring, old=False):
		is_mega = "mega" in str(namestring).lower()

		name = str(poke_id)
		if old:
			name = f"old/{name}"
		if is_mega:
			name = f"{name}-mega"

		try:
			clip = await self.get_clip(f"url:http://dillerm.io/pokemon/cries/{name}.ogg", ctx)
			clip.volume = 0.1
			await self.play_clip(clip, ctx)
		except Http404Error:
			return False
		except AudioPlayerNotFoundError:
			return True

		return True

	@commands.command(aliases=["pokemon"])
	async def pokedex(self, ctx, *, pokemon):
		"""Looks up information about the indicated pokemon

		Pokemon should be specified using either their name or id number

		Clicking on the pokemon's name will bring you to their wiki page

		Example:
		`{cmdpfx}pokedex charizard`"""
		with ctx.channel.typing():
			data, species_data = await self.get_pokemon_data(pokemon)

		types = []
		for t in sorted(data["types"], key=lambda t: t["slot"]):
			types.append(self.poke_type(t["type"]["name"]))

		flavor_text = localize(species_data["flavor_text_entries"], "flavor_text")
		flavor_text = flavor_text.replace("\n", " ")

		embed = discord.Embed(description=flavor_text, color=poke_color(species_data["color"]["name"]))
		embed.title = data["localized_name"] + f" #{data['id']}"
		embed.url = data["wiki_url"]

		embed.set_thumbnail(url=data["sprites"]["front_default"])

		embed.add_field(name=f"Type{'s' if len(types) > 1 else ''}", value=f"{''.join(types)}")
		if species_data.get("habitat"):
			embed.add_field(name="Habitat", value=f"{species_data['habitat']['name']}")
		embed.add_field(name="Weight", value=f"{data['weight'] / 10} kg")
		embed.add_field(name="Height", value=f"{data['height'] / 10} m")

		# fails silently if there's no cry for this pokemon
		await self.play_pokecry(ctx, data["id"], pokemon)

		await ctx.send(embed=embed)

	@commands.command()
	async def shiny(self, ctx, *, pokemon):
		"""Gets the shiny version of this pokemon

		pokemon should be specified using either their name or id number

		Example:
		`{cmdpfx}shiny charizard`"""

		# Sanitize input first
		with ctx.channel.typing():
			data, species_data = await self.get_pokemon_data(pokemon)

		if not data["sprites"].get("front_shiny"):
			await ctx.send("This pokemon doesn't have a shiny version")

		embed = discord.Embed(color=poke_color(species_data["color"]["name"]))
		embed.set_image(url=data["sprites"].get("front_shiny"))

		# fails silently if there's no cry for this pokemon
		await self.play_pokecry(ctx, data["id"], pokemon)

		await ctx.send(embed=embed)

	@commands.command(aliases=["cry"])
	async def pokecry(self, ctx, *, pokemon):
		"""Plays the pokemon's sound effect

		Audio files for these pokemon cries were gotten from [Veekun](https://veekun.com/dex/downloads). Veekun does not have the cries for Generation VII yet, so I won't be able to play those.

		Most pokemon have a new cry and an old cry. To get the old cry instead of the new one, add 'old' to the end of your command (see example below.)

		**Example:**
		`{cmdpfx}pokecry pikachu`
		`{cmdpfx}pokecry bulbasaur old`"""
		words = pokemon.split(" ")
		old = "old" in words
		if old:
			words.remove("old")
		pokemon = " ".join(words)

		data, species_data = await self.get_pokemon_data(pokemon)

		if data["id"] > 721:
			raise UserError("Sorry, I don't have the cries for pokemon in Generation VII yet")

		success = await self.play_pokecry(ctx, data["id"], pokemon, old=old)

		if not success:
			raise UserError(f"Couldn't find the cry for {data['localized_name']}")

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


<commit_before>import requests
from keen import exceptions

__author__ = 'dkador'

class KeenApi(object):
    base_url = "https://api.keen.io"
    api_version = "2.0"

    def __init__(self, project_id, auth_token, base_url=None,
                 api_version=None):
        super(KeenApi, self).__init__()
        self.project_id = project_id
        self.auth_token = auth_token
        if base_url:
            self.base_url = base_url
        if api_version:
            self.api_version = api_version

    def post_event(self, event):
        url = "{}/{}/projects/{}/{}".format(self.base_url, self.api_version,
                                            self.project_id,
                                            event.collection_name)
        headers = {"Authorization": self.auth_token,
                   "Content-Type": "application/json"}
        payload = event.to_json()
        response = requests.post(url, data=payload, headers=headers)
        if response.status_code != 201:
            error = response.json
            raise exceptions.KeenApiError(error)<commit_msg>add docstrings to api<commit_after>import requests
from keen import exceptions

__author__ = 'dkador'

class KeenApi(object):
    """
    Responsible for communicating with the Keen API. Used by multiple
    persistence strategies or async processing.
    """

    # the default base URL of the Keen API
    base_url = "https://api.keen.io"
    # the default version of the Keen API
    api_version = "2.0"

    def __init__(self, project_id, auth_token, base_url=None,
                 api_version=None):
        """ Initializes a KeenApi object

        :param project_id: the Keen project ID
        :param auth_token: the Keen authorization token
        :param base_url: optional, set this to override where API requests
        are sent
        :param api_version: string, optional, set this to override what API
        version is used
        """
        super(KeenApi, self).__init__()
        self.project_id = project_id
        self.auth_token = auth_token
        if base_url:
            self.base_url = base_url
        if api_version:
            self.api_version = api_version

    def post_event(self, event):
        """ Posts a single event to the Keen API.

        :param event: an Event to upload
        """
        url = "{}/{}/projects/{}/{}".format(self.base_url, self.api_version,
                                            self.project_id,
                                            event.collection_name)
        headers = {"Authorization": self.auth_token,
                   "Content-Type": "application/json"}
        payload = event.to_json()
        response = requests.post(url, data=payload, headers=headers)
        if response.status_code != 201:
            error = response.json
            raise exceptions.KeenApiError(error)<|endoftext|>


<commit_before>import pandas
import protocol
import numpy as np
import requests

from bokeh.properties import (HasProps, MetaHasProps,
        Any, Dict, Enum, Float, Instance, Int, List, String,
        Color, Pattern, Percent, Size, Bool)

#loading dependencies
import bokeh.objects
import bokeh.glyphs

from bokeh.objects import PlotObject, Plot, ColumnDataSource
from bokeh.session import PlotContext, PlotList

class IPythonRemoteData(PlotObject):
    host  = String("localhost")
    port = Int(10020)
    varname = String()
    computed_columns = List()

    #hack... we're just using this field right now to trigger events
    selected = Int(0)
    data = Int(0)

    def setselect(self, select, transform):

        remotedata = self
        url = "http://%s:%s/array/%s/setselect" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        data = transform
        data['selected'] = select
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def search(self, search):
        remotedata = self
        url = "http://%s:%s/array/%s/search" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        requests.post(url, data=search)
        self.selected += 1

    def select(self, select, transform):

        remotedata = self
        url = "http://%s:%s/array/%s/select" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        data = transform
        data['selected'] = select
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def deselect(self, deselect, transform):
        remotedata = self
        url = "http://%s:%s/array/%s/deselect" % (remotedata.host,
                                                  remotedata.port,
                                                  remotedata.varname)
        data = transform
        data['selected'] = deselect
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def get_data(self, transform):
        remotedata = self
        url = "http://%s:%s/array/%s" % (remotedata.host,
                                         remotedata.port,
                                         remotedata.varname)
        data = requests.get(url, data=protocol.serialize_json(transform)).json()
        return data

    def set_computed_columns(self, computed_columns):

        remotedata = self
        url = "http://%s:%s/array/%s/computed" % (remotedata.host,
                                                  remotedata.port,
                                                  remotedata.varname)
        data = requests.get(
            url,
            data=protocol.serialize_json(computed_columns)).json()
        self.computed_columns = computed_columns
        self.data += 1
        return data

class PandasPlotSource(ColumnDataSource):
    source = Instance(has_ref=True)

    def __init__(self, *args, **kwargs):
        super(PandasPlotSource, self).__init__(*args, **kwargs)

    def setup_events(self):
        self.on_change('selected', self, 'selection_callback')
        self.source.on_change('selected', self, 'get_data')
        self.source.on_change('data', self, 'get_data')
        self.source.on_change('computed_columns', self, 'get_data')
        if not self.data:
            self.get_data()

    def selection_callback(self, obj=None, attrname=None, old=None, new=None):
        self.setselect(self.selected)

    def transform(self):
        return {}

    def setselect(self, select):
        self.source.setselect(select, self.transform())
        self.get_data()

    def select(self, select):
        self.source.select(select, self.transform())
        self.get_data()

    def deselect(self, deselect):
        self.source.deselect(deselect, self.transform())
        self.get_data()

    def get_data(self, obj=None, attrname=None, old=None, new=None):
        data = self.source.get_data(self.transform())
        #ugly:
        self._selected =  np.nonzero(data['data']['_selected'])[0]
        self.maxlength = data.pop('maxlength')
        self.totallength = data.pop('totallength')
        self.column_names = data['column_names']
        self.data = data['data']

class PandasPivotTable(PlotObject):
    source = Instance(has_ref=True)
    sort = List()
    group = List()
    offset = Int(default=0)
    length = Int(default=100)
    maxlength = Int()
    totallength = Int()
    precision = Dict()
    tabledata = Dict()
    filterselected = Bool(default=False)
    def setup_events(self):
        self.on_change('sort', self, 'get_data')
        self.on_change('group', self, 'get_data')
        self.on_change('length', self, 'get_data')
        self.on_change('offset', self, 'get_data')
        self.on_change('precision', self, 'get_data')
        self.on_change('filterselected', self, 'get_data')
        self.source.on_change('selected', self, 'get_data')
        self.source.on_change('data', self, 'get_data')
        self.source.on_change('computed_columns', self, 'get_data')
        if not self.tabledata:
            self.get_data()

    def format_data(self, jsondata):
        """inplace manipulation of jsondata
        """
        precision = self.precision
        for colname, data in jsondata.iteritems():
            for idx, val in enumerate(data):
                if isinstance(val, float):
                    data[idx] = "%%.%df" % precision.get(colname,2)%data[idx]

    def transform(self):
        return dict(sort=self.sort,
                    group=self.group,
                    offset=self.offset,
                    length=self.length,
                    filterselected=self.filterselected,
                    )

    def setselect(self, select):
        self.source.setselect(select, self.transform())
        self.get_data()

    def select(self, select):
        self.source.select(select, self.transform())
        self.get_data()

    def deselect(self, deselect):
        self.source.deselect(deselect, self.transform())
        self.get_data()

    def get_data(self, obj=None, attrname=None, old=None, new=None):
        print 'get data table'
        data = self.source.get_data(self.transform())
        print data['data']['_selected']
        self.maxlength = data.pop('maxlength')
        self.totallength = data.pop('totallength')
        self.format_data(data['data'])
        self.tabledata = data

<commit_msg>dont' convert counts/selected to strings<commit_after>import pandas
import protocol
import numpy as np
import requests

from bokeh.properties import (HasProps, MetaHasProps,
        Any, Dict, Enum, Float, Instance, Int, List, String,
        Color, Pattern, Percent, Size, Bool)

#loading dependencies
import bokeh.objects
import bokeh.glyphs

from bokeh.objects import PlotObject, Plot, ColumnDataSource
from bokeh.session import PlotContext, PlotList

class IPythonRemoteData(PlotObject):
    host  = String("localhost")
    port = Int(10020)
    varname = String()
    computed_columns = List()

    #hack... we're just using this field right now to trigger events
    selected = Int(0)
    data = Int(0)

    def setselect(self, select, transform):

        remotedata = self
        url = "http://%s:%s/array/%s/setselect" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        data = transform
        data['selected'] = select
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def search(self, search):
        remotedata = self
        url = "http://%s:%s/array/%s/search" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        requests.post(url, data=search)
        self.selected += 1

    def select(self, select, transform):

        remotedata = self
        url = "http://%s:%s/array/%s/select" % (remotedata.host,
                                                remotedata.port,
                                                remotedata.varname)
        data = transform
        data['selected'] = select
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def deselect(self, deselect, transform):
        remotedata = self
        url = "http://%s:%s/array/%s/deselect" % (remotedata.host,
                                                  remotedata.port,
                                                  remotedata.varname)
        data = transform
        data['selected'] = deselect
        requests.post(url, data=protocol.serialize_json(data))
        self.selected += 1

    def get_data(self, transform):
        remotedata = self
        url = "http://%s:%s/array/%s" % (remotedata.host,
                                         remotedata.port,
                                         remotedata.varname)
        data = requests.get(url, data=protocol.serialize_json(transform)).json()
        return data

    def set_computed_columns(self, computed_columns):

        remotedata = self
        url = "http://%s:%s/array/%s/computed" % (remotedata.host,
                                                  remotedata.port,
                                                  remotedata.varname)
        data = requests.get(
            url,
            data=protocol.serialize_json(computed_columns)).json()
        self.computed_columns = computed_columns
        self.data += 1
        return data

class PandasPlotSource(ColumnDataSource):
    source = Instance(has_ref=True)

    def __init__(self, *args, **kwargs):
        super(PandasPlotSource, self).__init__(*args, **kwargs)

    def setup_events(self):
        self.on_change('selected', self, 'selection_callback')
        self.source.on_change('selected', self, 'get_data')
        self.source.on_change('data', self, 'get_data')
        self.source.on_change('computed_columns', self, 'get_data')
        if not self.data:
            self.get_data()

    def selection_callback(self, obj=None, attrname=None, old=None, new=None):
        self.setselect(self.selected)

    def transform(self):
        return {}

    def setselect(self, select):
        self.source.setselect(select, self.transform())
        self.get_data()

    def select(self, select):
        self.source.select(select, self.transform())
        self.get_data()

    def deselect(self, deselect):
        self.source.deselect(deselect, self.transform())
        self.get_data()

    def get_data(self, obj=None, attrname=None, old=None, new=None):
        data = self.source.get_data(self.transform())
        #ugly:
        self._selected =  np.nonzero(data['data']['_selected'])[0]
        self.maxlength = data.pop('maxlength')
        self.totallength = data.pop('totallength')
        self.column_names = data['column_names']
        self.data = data['data']

class PandasPivotTable(PlotObject):
    source = Instance(has_ref=True)
    sort = List()
    group = List()
    offset = Int(default=0)
    length = Int(default=100)
    maxlength = Int()
    totallength = Int()
    precision = Dict()
    tabledata = Dict()
    filterselected = Bool(default=False)
    def setup_events(self):
        self.on_change('sort', self, 'get_data')
        self.on_change('group', self, 'get_data')
        self.on_change('length', self, 'get_data')
        self.on_change('offset', self, 'get_data')
        self.on_change('precision', self, 'get_data')
        self.on_change('filterselected', self, 'get_data')
        self.source.on_change('selected', self, 'get_data')
        self.source.on_change('data', self, 'get_data')
        self.source.on_change('computed_columns', self, 'get_data')
        if not self.tabledata:
            self.get_data()

    def format_data(self, jsondata):
        """inplace manipulation of jsondata
        """
        precision = self.precision
        for colname, data in jsondata.iteritems():
            if colname == '_selected' or colname == '_counts':
                continue
            for idx, val in enumerate(data):
                if isinstance(val, float):
                    data[idx] = "%%.%df" % precision.get(colname,2)%data[idx]

    def transform(self):
        return dict(sort=self.sort,
                    group=self.group,
                    offset=self.offset,
                    length=self.length,
                    filterselected=self.filterselected,
                    )

    def setselect(self, select):
        self.source.setselect(select, self.transform())
        self.get_data()

    def select(self, select):
        self.source.select(select, self.transform())
        self.get_data()

    def deselect(self, deselect):
        self.source.deselect(deselect, self.transform())
        self.get_data()

    def get_data(self, obj=None, attrname=None, old=None, new=None):
        data = self.source.get_data(self.transform())
        print data['data']['_selected']
        self.maxlength = data.pop('maxlength')
        self.totallength = data.pop('totallength')
        self.format_data(data['data'])
        self.tabledata = data

<|endoftext|>


<commit_before>import numpy, sys, time, pylab, math, re
from astropy.io import fits as pyfits
from pylab import *
from matplotlib import *
import matplotlib.mlab as mlab
from math import *
from numpy import *
import kepio, kepmsg, kepkey, kepfit, kepfunc, kepstat, kepfourier

def keptrial(infile,outfile,datacol,errcol,fmin,fmax,nfreq,method,
             ntrials,plot,clobber,verbose,logfile,status,cmdLine=False):

# startup parameters

    status = 0
    labelsize = 24
    ticksize = 16
    xsize = 18
    ysize = 6
    lcolor = '#0000ff'
    lwidth = 1.0
    fcolor = '#ffff00'
    falpha = 0.2

# log the call

    hashline = '----------------------------------------------------------------------------'
    kepmsg.log(logfile,hashline,verbose)
    call = 'KEPTRIAL -- '
    call += 'infile='+infile+' '
    call += 'outfile='+outfile+' '
    call += 'datacol='+datacol+' '
    call += 'errcol='+errcol+' '
    call += 'fmin='+str(fmin)+' '
    call += 'fmax='+str(fmax)+' '
    call += 'nfreq='+str(nfreq)+' '
    call += 'method='+method+' '
    call += 'ntrials='+str(ntrials)+' '
    plotit = 'n'
    if (plot): plotit = 'y'
    call += 'plot='+plotit+ ' '
    overwrite = 'n'
    if (clobber): overwrite = 'y'
    call += 'clobber='+overwrite+ ' '
    chatter = 'n'
    if (verbose): chatter = 'y'
    call += 'verbose='+chatter+' '
    call += 'logfile='+logfile
    kepmsg.log(logfile,call+'\n',verbose)

# start time

    kepmsg.clock('KEPTRIAL started at',logfile,verbose)

# test log file

    logfile = kepmsg.test(logfile)

# clobber output file

    if clobber: status = kepio.clobber(outfile,logfile,verbose)
    if kepio.fileexists(outfile):
	    message = 'ERROR -- KEPTRIAL: ' + outfile + ' exists. Use clobber=yes'
	    kepmsg.err(logfile,message,verbose)
	    status = 1

# open input file

    if status == 0:
        instr, status = kepio.openfits(infile,'readonly',logfile,verbose)

# fudge non-compliant FITS keywords with no values

    if status == 0:
        instr = kepkey.emptykeys(instr,file,logfile,verbose)

# input data

    if status == 0:
	try:
            barytime = instr[1].data.field('barytime')
	except:
            barytime, status = kepio.readfitscol(infile,instr[1].data,'time',logfile,verbose)
    if status == 0:
        signal, status = kepio.readfitscol(infile,instr[1].data,datacol,logfile,verbose)
    if status == 0:
        err, status = kepio.readfitscol(infile,instr[1].data,errcol,logfile,verbose)

# remove infinite data from time series

    if status == 0:
        try:
            nanclean = instr[1].header['NANCLEAN']
        except:
	    incols = [barytime, signal, err]
	    [barytime, signal, err] = kepstat.removeinfinlc(signal, incols)

# set up plot

    if status == 0:
        plotLatex = True
        try:
            params = {'backend': 'png',
                      'axes.linewidth': 2.5,
                      'axes.labelsize': labelsize,
                      'axes.font': 'sans-serif',
                      'axes.fontweight' : 'bold',
                      'text.fontsize': 12,
                      'legend.fontsize': 12,
                      'xtick.labelsize': ticksize,
                      'ytick.labelsize': ticksize}
            rcParams.update(params)
        except:
            print 'WARNING: install latex for scientific plotting'
            plotLatex = False

# frequency steps and Monte Carlo iterations

    if status == 0:
        deltaf = (fmax - fmin) / nfreq
        freq = []; pmax = []; trial = []
        for i in range(ntrials):
            trial.append(i+1)

# adjust data within the error bars

            work1 = kepstat.randarray(signal,err)

# determine FT power
            fr, power = kepfourier.ft(barytime,work1,fmin,fmax,deltaf,False)

# determine peak in FT

            pmax.append(-1.0e30)
            for j in range(len(fr)):
                if (power[j] > pmax[-1]):
                    pmax[-1] = power[j]
                    f1 = fr[j]
            freq.append(f1)

# plot stop-motion histogram

            pylab.ion()
	    pylab.figure(1,figsize=[7,10])
            clf()
	    pylab.axes([0.08,0.08,0.88,0.89])
            pylab.gca().xaxis.set_major_formatter(pylab.ScalarFormatter(useOffset=False))
            pylab.gca().yaxis.set_major_formatter(pylab.ScalarFormatter(useOffset=False))
            n,bins,patches = pylab.hist(freq,bins=nfreq,range=[fmin,fmax],
                                        align='mid',rwidth=1,ec='#0000ff',
                                        fc='#ffff00',lw=2)

# fit normal distribution to histogram

            x = zeros(len(bins))
            for j in range(1,len(bins)):
                x[j] = (bins[j] + bins[j-1]) / 2
            pinit = numpy.array([float(i),freq[-1],deltaf])
            if i > 3:
                n = array(n,dtype='float32')
                coeffs, errors, covar, sigma, chi2, dof, fit, plotx, ploty, status = \
                    kepfit.leastsquare('gauss',pinit,x[1:],n,None,logfile,verbose)
                fitfunc = kepfunc.gauss()
                f = arange(fmin,fmax,(fmax-fmin)/100)
                fit = fitfunc(coeffs,f)
                pylab.plot(f,fit,'r-',linewidth=2)
            if plotLatex:
                xlabel(r'Frequency (d$^{-1}$)', {'color' : 'k'})
            else:
                xlabel(r'Frequency (1/d)', {'color' : 'k'})
            ylabel('N', {'color' : 'k'})
            xlim(fmin,fmax)
	    grid()

# render plot

        if plot:
            if cmdLine:
                pylab.show()
            else:
                pylab.ion()
                pylab.plot([])
                pylab.ioff()

# period results

    if status == 0:
        p = 1.0 / coeffs[1]
        perr = p * coeffs[2] / coeffs[1]
        f1 = fmin; f2 = fmax
        gotbin = False
        for i in range(len(n)):
            if n[i] > 0 and not gotbin:
                f1 = bins[i]
                gotbin = True
        gotbin = False
        for i in range(len(n)-1,0,-1):
            if n[i] > 0 and not gotbin:
                f2 = bins[i+1]
                gotbin = True
        powave, powstdev = kepstat.stdev(pmax)

# print result

    if status == 0:
        print '              best period: %.10f days (%.7f min)' % (p, p * 1440.0)
        print '     1-sigma period error: %.10f days (%.7f min)' % (perr, perr * 1440.0)
        print '             search range: %.10f - %.10f days  ' % (1.0 / fmax, 1.0 / fmin)
        print '    100%% confidence range: %.10f - %.10f days  ' % (1.0 / f2, 1.0 / f1)
#        print '     detection confidence: %.2f sigma' % (powave / powstdev)
        print '         number of trials: %d' % ntrials
        print ' number of frequency bins: %d' % nfreq

# history keyword in output file

    if status == 0:
	    status = kepkey.history(call,instr[0],outfile,logfile,verbose)

## write output file

    if status == 0:
        col1 = Column(name='TRIAL',format='J',array=trial)
        col2 = Column(name='FREQUENCY',format='E',unit='1/day',array=freq)
        col3 = Column(name='POWER',format='E',array=pmax)
        cols = ColDefs([col1,col2,col3])
        instr.append(new_table(cols))
        try:
            instr[-1].header.update('EXTNAME','TRIALS','Extension name')
        except:
            status = 1
        try:
            instr[-1].header.update('SEARCHR1',1.0 / fmax,'Search range lower bound (days)')
        except:
            status = 1
        try:
            instr[-1].header.update('SEARCHR2',1.0 / fmin,'Search range upper bound (days)')
        except:
            status = 1
        try:
            instr[-1].header.update('NFREQ',nfreq,'Number of frequency bins')
        except:
            status = 1
        try:
            instr[-1].header.update('PERIOD',p,'Best period (days)')
        except:
            status = 1
        try:
            instr[-1].header.update('PERIODE',perr,'1-sigma period error (days)')
        except:
            status = 1
#        instr[-1].header.update('DETNCONF',powave/powstdev,'Detection significance (sigma)')
        try:
            instr[-1].header.update('CONFIDR1',1.0 / f2,'Trial confidence lower bound (days)')
        except:
            status = 1
        try:
            instr[-1].header.update('CONFIDR2',1.0 / f1,'Trial confidence upper bound (days)')
        except:
            status = 1
        try:
            instr[-1].header.update('NTRIALS',ntrials,'Number of trials')
        except:
            status = 1
        instr.writeto(outfile)

# close input file

    if status == 0:
	    status = kepio.closefits(instr,logfile,verbose)

## end time

    if (status == 0):
	    message = 'KEPTRAIL completed at'
    else:
	    message = '\nKEPTRIAL aborted at'
    kepmsg.clock(message,logfile,verbose)

# main
if '--shell' in sys.argv:
    import argparse

    parser = argparse.ArgumentParser(description='Calculate best period and error estimate from Fourier transform')
    parser.add_argument('--shell', action='store_true', help='Are we running from the shell?')
    parser.add_argument('infile', help='Name of input file', type=str)

    parser.add_argument('outfile', help='Name of FITS file to output', type=str)

    parser.add_argument('--datacol', default='SAP_FLUX', help='Name of data column', type=str)
    parser.add_argument('--errcol', default='SAP_FLUX_ERR', help='Name of data error column', type=str)

    parser.add_argument('--fmin', default=0.1, help='Minimum search frequency [1/day]', type=float)
    parser.add_argument('--fmax', default=50., help='Minimum search frequency [1/day]', type=float)
    parser.add_argument('--nfreq', default=100, help='Number of frequency intervals', type=int)

    parser.add_argument('--method', default='ft', help='Frequency search method', type=int, choices=['ft'])
    parser.add_argument('--ntrials', default=1000, help='Number of search trials', type=int)

    parser.add_argument('--plot', action='store_true', help='Plot result?')

    parser.add_argument('--clobber', action='store_true', help='Overwrite output file?')
    parser.add_argument('--verbose', action='store_true', help='Write to a log file?')
    parser.add_argument('--logfile', '-l', help='Name of ascii log file', default='kepcotrend.log', dest='logfile', type=str)
    parser.add_argument('--status', '-e', help='Exit status (0=good)', default=0, dest='status', type=int)

    args = parser.parse_args()

    cmdLine=True

    keptrial(args.infile,args.outfile,args.datacol,args.errcol,args.fmin,args.fmax,args.nfreq,args.method,
             args.ntrials,args.plot,args.clobber,args.verbose,args.logfile,args.status, cmdLine)

else:
    from pyraf import iraf
    parfile = iraf.osfn("kepler$keptrial.par")
    t = iraf.IrafTaskFactory(taskname="keptrial", value=parfile, function=keptrial)
<commit_msg>fix keptrial<commit_after>import sys, time, math, re
from astropy.io import fits as pyfits
from matplotlib import pyplot as plt
import numpy as np
import kepio, kepmsg, kepkey, kepfit, kepfunc, kepstat, kepfourier

def keptrial(infile,outfile,datacol,errcol,fmin,fmax,nfreq,method,
             ntrials,plot,clobber,verbose,logfile,status,cmdLine=False):

# startup parameters

    status = 0
    labelsize = 24
    ticksize = 16
    xsize = 18
    ysize = 6
    lcolor = '#0000ff'
    lwidth = 1.0
    fcolor = '#ffff00'
    falpha = 0.2

# log the call

    hashline = '----------------------------------------------------------------------------'
    kepmsg.log(logfile,hashline,verbose)
    call = 'KEPTRIAL -- '
    call += 'infile='+infile+' '
    call += 'outfile='+outfile+' '
    call += 'datacol='+datacol+' '
    call += 'errcol='+errcol+' '
    call += 'fmin='+str(fmin)+' '
    call += 'fmax='+str(fmax)+' '
    call += 'nfreq='+str(nfreq)+' '
    call += 'method='+method+' '
    call += 'ntrials='+str(ntrials)+' '
    plotit = 'n'
    if (plot): plotit = 'y'
    call += 'plot='+plotit+ ' '
    overwrite = 'n'
    if (clobber): overwrite = 'y'
    call += 'clobber='+overwrite+ ' '
    chatter = 'n'
    if (verbose): chatter = 'y'
    call += 'verbose='+chatter+' '
    call += 'logfile='+logfile
    kepmsg.log(logfile,call+'\n',verbose)

# start time

    kepmsg.clock('KEPTRIAL started at',logfile,verbose)

# test log file

    logfile = kepmsg.test(logfile)

# clobber output file

    if clobber: status = kepio.clobber(outfile,logfile,verbose)
    if kepio.fileexists(outfile):
	    message = 'ERROR -- KEPTRIAL: ' + outfile + ' exists. Use clobber=yes'
	    kepmsg.err(logfile,message,verbose)
	    status = 1

# open input file

    if status == 0:
        instr, status = kepio.openfits(infile,'readonly',logfile,verbose)

# fudge non-compliant FITS keywords with no values

    if status == 0:
        instr = kepkey.emptykeys(instr,file,logfile,verbose)

# input data

    if status == 0:
	try:
            barytime = instr[1].data.field('barytime')
	except:
            barytime, status = kepio.readfitscol(infile,instr[1].data,'time',logfile,verbose)
    if status == 0:
        signal, status = kepio.readfitscol(infile,instr[1].data,datacol,logfile,verbose)
    if status == 0:
        err, status = kepio.readfitscol(infile,instr[1].data,errcol,logfile,verbose)

# remove infinite data from time series

    if status == 0:
        try:
            nanclean = instr[1].header['NANCLEAN']
        except:
	    incols = [barytime, signal, err]
	    [barytime, signal, err] = kepstat.removeinfinlc(signal, incols)

# frequency steps and Monte Carlo iterations

    if status == 0:
        deltaf = (fmax - fmin) / nfreq
        freq = []; pmax = []; trial = []
        for i in range(ntrials):
            trial.append(i+1)

# adjust data within the error bars

            work1 = kepstat.randarray(signal,err)

# determine FT power
            fr, power = kepfourier.ft(barytime,work1,fmin,fmax,deltaf,False)

# determine peak in FT

            pmax.append(-1.0e30)
            for j in range(len(fr)):
                if (power[j] > pmax[-1]):
                    pmax[-1] = power[j]
                    f1 = fr[j]
            freq.append(f1)

# plot stop-motion histogram

            plt.ion()
	    plt.figure(1,figsize=[7,10])
            plt.clf()
	    plt.axes([0.08,0.08,0.88,0.89])
            plt.gca().xaxis.set_major_formatter(plt.ScalarFormatter(useOffset=False))
            plt.gca().yaxis.set_major_formatter(plt.ScalarFormatter(useOffset=False))
            n,bins,patches = plt.hist(freq,bins=nfreq,range=[fmin,fmax],
                                      align='mid',rwidth=1,ec='#0000ff',
                                      fc='#ffff00',lw=2)

# fit normal distribution to histogram

            x = np.zeros(len(bins))
            for j in range(1,len(bins)):
                x[j] = (bins[j] + bins[j-1]) / 2
            pinit = np.array([float(i),freq[-1],deltaf])
            if i > 3:
                n = np.array(n,dtype='float32')
                coeffs, errors, covar, sigma, chi2, dof, fit, plotx, ploty, status = \
                    kepfit.leastsquare('gauss',pinit,x[1:],n,None,logfile,verbose)
                fitfunc = kepfunc.gauss()
                f = np.arange(fmin,fmax,(fmax-fmin)/100)
                fit = fitfunc(coeffs,f)
                plt.plot(f,fit,'r-',linewidth=2)
                plt.xlabel(r'Frequency (1/d)', {'color' : 'k'})
            plt.ylabel('N', {'color' : 'k'})
            plt.xlim(fmin,fmax)
	    plt.grid()

# render plot

        if plot:
            plt.show()
# period results

    if status == 0:
        p = 1.0 / coeffs[1]
        perr = p * coeffs[2] / coeffs[1]
        f1 = fmin; f2 = fmax
        gotbin = False
        for i in range(len(n)):
            if n[i] > 0 and not gotbin:
                f1 = bins[i]
                gotbin = True
        gotbin = False
        for i in range(len(n)-1,0,-1):
            if n[i] > 0 and not gotbin:
                f2 = bins[i+1]
                gotbin = True
        powave, powstdev = kepstat.stdev(pmax)

# print result

    if status == 0:
        print '              best period: %.10f days (%.7f min)' % (p, p * 1440.0)
        print '     1-sigma period error: %.10f days (%.7f min)' % (perr, perr * 1440.0)
        print '             search range: %.10f - %.10f days  ' % (1.0 / fmax, 1.0 / fmin)
        print '    100%% confidence range: %.10f - %.10f days  ' % (1.0 / f2, 1.0 / f1)
#        print '     detection confidence: %.2f sigma' % (powave / powstdev)
        print '         number of trials: %d' % ntrials
        print ' number of frequency bins: %d' % nfreq

# history keyword in output file

    if status == 0:
	    status = kepkey.history(call,instr[0],outfile,logfile,verbose)

## write output file

    if status == 0:
        col1 = pyfits.Column(name='TRIAL',format='J',array=trial)
        col2 = pyfits.Column(name='FREQUENCY',format='E',unit='1/day',array=freq)
        col3 = pyfits.Column(name='POWER',format='E',array=pmax)
        cols = pyfits.ColDefs([col1,col2,col3])
        instr.append(pyfits.BinTableHDU.from_columns(cols))
        try:
            instr[-1].header['EXTNAME' ] = ('TRIALS','Extension name')
        except:
            status = 1
        try:
            instr[-1].header['SEARCHR1'] = (1.0 / fmax,'Search range lower bound (days)')
        except:
            status = 1
        try:
            instr[-1].header['SEARCHR2'] = (1.0 / fmin,'Search range upper bound (days)')
        except:
            status = 1
        try:
            instr[-1].header['NFREQ'   ] = (nfreq,'Number of frequency bins')
        except:
            status = 1
        try:
            instr[-1].header['PERIOD'  ] = (p,'Best period (days)')
        except:
            status = 1
        try:
            instr[-1].header['PERIODE' ] = (perr,'1-sigma period error (days)')
        except:
            status = 1
        try:
            instr[-1].header['CONFIDR1'] = (1.0 / f2,'Trial confidence lower bound (days)')
        except:
            status = 1
        try:
            instr[-1].header['CONFIDR2'] = (1.0 / f1,'Trial confidence upper bound (days)')
        except:
            status = 1
        try:
            instr[-1].header['NTRIALS' ] = (ntrials,'Number of trials')
        except:
            status = 1
        instr.writeto(outfile)

# close input file

    if status == 0:
	    status = kepio.closefits(instr,logfile,verbose)

## end time

    if (status == 0):
	    message = 'KEPTRAIL completed at'
    else:
	    message = '\nKEPTRIAL aborted at'
    kepmsg.clock(message,logfile,verbose)

# main
if '--shell' in sys.argv:
    import argparse
    parser = argparse.ArgumentParser(description='Calculate best period and error estimate from Fourier transform')
    parser.add_argument('--shell', action='store_true', help='Are we running from the shell?')
    parser.add_argument('infile', help='Name of input file', type=str)

    parser.add_argument('outfile', help='Name of FITS file to output', type=str)

    parser.add_argument('--datacol', default='SAP_FLUX', help='Name of data column', type=str)
    parser.add_argument('--errcol', default='SAP_FLUX_ERR', help='Name of data error column', type=str)

    parser.add_argument('--fmin', default=0.1, help='Minimum search frequency [1/day]', type=float)
    parser.add_argument('--fmax', default=50., help='Minimum search frequency [1/day]', type=float)
    parser.add_argument('--nfreq', default=100, help='Number of frequency intervals', type=int)

    parser.add_argument('--method', default='ft', help='Frequency search method', type=int, choices=['ft'])
    parser.add_argument('--ntrials', default=1000, help='Number of search trials', type=int)

    parser.add_argument('--plot', action='store_true', help='Plot result?')

    parser.add_argument('--clobber', action='store_true', help='Overwrite output file?')
    parser.add_argument('--verbose', action='store_true', help='Write to a log file?')
    parser.add_argument('--logfile', '-l', help='Name of ascii log file',
                        default='keptrial.log', dest='logfile', type=str)
    parser.add_argument('--status', '-e', help='Exit status (0=good)', default=0, dest='status', type=int)

    args = parser.parse_args()
    cmdLine=True

    keptrial(args.infile,args.outfile,args.datacol,args.errcol,args.fmin,args.fmax,args.nfreq,args.method,
             args.ntrials,args.plot,args.clobber,args.verbose,args.logfile,args.status, cmdLine)

else:
    from pyraf import iraf
    parfile = iraf.osfn("kepler$keptrial.par")
    t = iraf.IrafTaskFactory(taskname="keptrial", value=parfile, function=keptrial)
<|endoftext|>


<commit_before>import discord
from discord.ext import commands
from ext.colors import ColorNames
from PIL import Image
import inspect
import io

class Utility:
    '''Useful commands to make your life easier'''

    def __init__(self, bot):
        self.bot = bot

    @commands.command(name='logout')
    async def _logout(self, ctx):
        '''
        Shuts down the selfbot,
        equivalent to a restart if you are hosting on heroku
        '''
        await ctx.send('`Selfbot Logging out...`')
        await self.bot.logout()

    @commands.command()
    async def dcolor(self, ctx, *, url):
        '''A command that gets the dominant color of an image url'''
        await ctx.message.delete()
        color = await ctx.get_dominant_color(url)
        string_col = ColorNames.color_name(str(color))
        info = f'`{str(color)}`\n`{color.to_rgb()}`\n`{str(string_col)}`'
        em = discord.Embed(color=color, title='Dominant Color', description=info)
        em.set_thumbnail(url=url)
        em.set_footer(text='Selfbot made by SharpBit | Powered by cr-api',
                      icon_url='http://cr-api.com/static/img/branding/cr-api-logo.png')
        file = io.BytesIO()
        Image.new('RGB', (200, 90), color.to_rgb()).save(file, format='PNG')
        file.seek(0)
        em.set_image(url="attachment://color.png")
        await ctx.send(file=discord.File(file, 'color.png'), embed=em)

    @commands.command()
    async def tinyurl(self, ctx, *, link: str):
        '''Makes a link shorter using the tinyurl api'''
        await ctx.message.delete()
        url = 'http://tinyurl.com/api-create.php?url=' + link
        async with ctx.session.get(url) as resp:
            new = await resp.text()
        emb = discord.Embed(color=await ctx.get_dominant_color(ctx.author.avatar_url))
        emb.add_field(name="Original Link", value=link, inline=False)
        emb.add_field(name="Shortened Link", value=new, inline=False)
        emb.set_footer(text='Selfbot made by SharpBit | Powered by cr-api',
                       icon_url='http://cr-api.com/static/img/branding/cr-api-logo.png')
        await ctx.send(embed=emb)

    @commands.command()
    async def hastebin(self, ctx, *, code):
        '''Hastebin-ify your code!'''
        async with ctx.session.post("https://hastebin.com/documents", data=code) as resp:
            data = await resp.json()
            key = data['key']
        await ctx.message.edit(content=f"Hastebin-ified! <https://hastebin.com/{key}.py>")

    @commands.command()
    async def source(self, ctx, *, command):
        '''See the source code for any command.'''
        source = str(inspect.getsource(self.bot.get_command(command).callback))
        fmt = '```py\n' + source.replace('`', '\u200b`') + '\n```'
        if len(fmt) > 2000:
            async with ctx.session.post("https://hastebin.com/documents", data=source) as resp:
                data = await resp.json()
            key = data['key']
            return await ctx.send(f'Command source: <https://hastebin.com/{key}.py>')
        else:
            return await ctx.send(fmt)

def setup(bot):
    bot.add_cog(Utility(bot))
<commit_msg>Custom help message?<commit_after>import discord
from discord.ext import commands
from ext.colors import ColorNames
from PIL import Image
import inspect
import io

class Utility:
    '''Useful commands to make your life easier'''

    def __init__(self, bot):
        self.bot = bot

    @commands.command(name='logout')
    async def _logout(self, ctx):
        '''
        Shuts down the selfbot,
        equivalent to a restart if you are hosting on heroku
        '''
        await ctx.send('`Selfbot Logging out...`')
        await self.bot.logout()

    @commands.command(name='help')
    async def new_help_command(self, ctx, *commands: str):
        """Shows this message."""
        destination = ctx.message.author if self.bot.pm_help else ctx.message.channel

        def repl(obj):
            return self.bot._mentions_transforms.get(obj.group(0), '')

        # help by itself just lists our own commands.
        if len(commands) == 0:
            pages = await self.bot.formatter.format_help_for(ctx, self.bot)
        elif len(commands) == 1:
            # try to see if it is a cog name
            name = self.bot._mention_pattern.sub(repl, commands[0])
            command = None
            if name in self.bot.cogs:
                command = self.bot.cogs[name]
            else:
                command = self.bot.all_commands.get(name)
                if command is None:
                    await destination.send(self.bot.command_not_found.format(name))
                    return

            pages = await self.bot.formatter.format_help_for(ctx, command)
        else:
            name = self.bot._mention_pattern.sub(repl, commands[0])
            command = self.bot.all_commands.get(name)
            if command is None:
                await destination.send(self.bot.command_not_found.format(name))
                return

            for key in commands[1:]:
                try:
                    key = self.bot._mention_pattern.sub(repl, key)
                    command = command.all_commands.get(key)
                    if command is None:
                        await destination.send(self.bot.command_not_found.format(key))
                        return
                except AttributeError:
                    await destination.send(self.bot.command_has_no_subcommands.format(command, key))
                    return

            pages = await self.bot.formatter.format_help_for(ctx, command)

        if self.bot.pm_help is None:
            characters = sum(map(lambda l: len(l), pages))
            # modify destination based on length of pages.
            if characters > 1000:
                destination = ctx.message.author

        color = await ctx.get_dominant_color(ctx.author.avatar_url)

        for embed in pages:
            em = discord.Embed(title='Command Help', color=color)
            embed += f'\n\nType {ctx.prefix}help command for more info on a command.\nYou can also type {ctx.prefix}help category for more info on a category.
            em.description = embed
            try:
                await ctx.send(embed=em)
            except discord.HTTPException:
                em_list = await embedtobox.etb(embed)
                for page in em_list:
                    await ctx.send(page)

    @commands.command()
    async def dcolor(self, ctx, *, url):
        '''A command that gets the dominant color of an image url'''
        await ctx.message.delete()
        color = await ctx.get_dominant_color(url)
        string_col = ColorNames.color_name(str(color))
        info = f'`{str(color)}`\n`{color.to_rgb()}`\n`{str(string_col)}`'
        em = discord.Embed(color=color, title='Dominant Color', description=info)
        em.set_thumbnail(url=url)
        em.set_footer(text='Selfbot made by SharpBit | Powered by cr-api',
                      icon_url='http://cr-api.com/static/img/branding/cr-api-logo.png')
        file = io.BytesIO()
        Image.new('RGB', (200, 90), color.to_rgb()).save(file, format='PNG')
        file.seek(0)
        em.set_image(url="attachment://color.png")
        await ctx.send(file=discord.File(file, 'color.png'), embed=em)

    @commands.command()
    async def tinyurl(self, ctx, *, link: str):
        '''Makes a link shorter using the tinyurl api'''
        await ctx.message.delete()
        url = 'http://tinyurl.com/api-create.php?url=' + link
        async with ctx.session.get(url) as resp:
            new = await resp.text()
        emb = discord.Embed(color=await ctx.get_dominant_color(ctx.author.avatar_url))
        emb.add_field(name="Original Link", value=link, inline=False)
        emb.add_field(name="Shortened Link", value=new, inline=False)
        emb.set_footer(text='Selfbot made by SharpBit | Powered by cr-api',
                       icon_url='http://cr-api.com/static/img/branding/cr-api-logo.png')
        await ctx.send(embed=emb)

    @commands.command()
    async def hastebin(self, ctx, *, code):
        '''Hastebin-ify your code!'''
        async with ctx.session.post("https://hastebin.com/documents", data=code) as resp:
            data = await resp.json()
            key = data['key']
        await ctx.message.edit(content=f"Hastebin-ified! <https://hastebin.com/{key}.py>")

    @commands.command()
    async def source(self, ctx, *, command):
        '''See the source code for any command.'''
        source = str(inspect.getsource(self.bot.get_command(command).callback))
        fmt = '```py\n' + source.replace('`', '\u200b`') + '\n```'
        if len(fmt) > 2000:
            async with ctx.session.post("https://hastebin.com/documents", data=source) as resp:
                data = await resp.json()
            key = data['key']
            return await ctx.send(f'Command source: <https://hastebin.com/{key}.py>')
        else:
            return await ctx.send(fmt)

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