<commit_before># -*- coding: utf-8 -*-

"""
The MIT License (MIT)

Copyright (c) 2015-2020 Rapptz

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""

import inspect
import copy
from ._types import _BaseCommand

__all__ = (
    'CogMeta',
    'Cog',
)

class CogMeta(type):
    """A metaclass for defining a cog.

    Note that you should probably not use this directly. It is exposed
    purely for documentation purposes along with making custom metaclasses to intermix
    with other metaclasses such as the :class:`abc.ABCMeta` metaclass.

    For example, to create an abstract cog mixin class, the following would be done.

    .. code-block:: python3

        import abc

        class CogABCMeta(commands.CogMeta, abc.ABCMeta):
            pass

        class SomeMixin(metaclass=abc.ABCMeta):
            pass

        class SomeCogMixin(SomeMixin, commands.Cog, metaclass=CogABCMeta):
            pass

    .. note::

        When passing an attribute of a metaclass that is documented below, note
        that you must pass it as a keyword-only argument to the class creation
        like the following example:

        .. code-block:: python3

            class MyCog(commands.Cog, name='My Cog'):
                pass

    Attributes
    -----------
    name: :class:`str`
        The cog name. By default, it is the name of the class with no modification.
    command_attrs: :class:`dict`
        A list of attributes to apply to every command inside this cog. The dictionary
        is passed into the :class:`Command` options at ``__init__``.
        If you specify attributes inside the command attribute in the class, it will
        override the one specified inside this attribute. For example:

        .. code-block:: python3

            class MyCog(commands.Cog, command_attrs=dict(hidden=True)):
                @commands.command()
                async def foo(self, ctx):
                    pass # hidden -> True

                @commands.command(hidden=False)
                async def bar(self, ctx):
                    pass # hidden -> False
    """

    def __new__(cls, *args, **kwargs):
        name, bases, attrs = args
        attrs['__cog_name__'] = kwargs.pop('name', name)
        attrs['__cog_settings__'] = command_attrs = kwargs.pop('command_attrs', {})

        commands = {}
        listeners = {}
        no_bot_cog = 'Commands or listeners must not start with cog_ or bot_ (in method {0.__name__}.{1})'

        new_cls = super().__new__(cls, name, bases, attrs, **kwargs)
        for base in reversed(new_cls.__mro__):
            for elem, value in base.__dict__.items():
                if elem in commands:
                    del commands[elem]
                if elem in listeners:
                    del listeners[elem]

                is_static_method = isinstance(value, staticmethod)
                if is_static_method:
                    value = value.__func__
                if isinstance(value, _BaseCommand):
                    if is_static_method:
                        raise TypeError('Command in method {0}.{1!r} must not be staticmethod.'.format(base, elem))
                    if elem.startswith(('cog_', 'bot_')):
                        raise TypeError(no_bot_cog.format(base, elem))
                    commands[elem] = value
                elif inspect.iscoroutinefunction(value):
                    try:
                        is_listener = getattr(value, '__cog_listener__')
                    except AttributeError:
                        continue
                    else:
                        if elem.startswith(('cog_', 'bot_')):
                            raise TypeError(no_bot_cog.format(base, elem))
                        listeners[elem] = value

        new_cls.__cog_commands__ = list(commands.values()) # this will be copied in Cog.__new__

        listeners_as_list = []
        for listener in listeners.values():
            for listener_name in listener.__cog_listener_names__:
                # I use __name__ instead of just storing the value so I can inject
                # the self attribute when the time comes to add them to the bot
                listeners_as_list.append((listener_name, listener.__name__))

        new_cls.__cog_listeners__ = listeners_as_list
        return new_cls

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

    @classmethod
    def qualified_name(cls):
        return cls.__cog_name__

def _cog_special_method(func):
    func.__cog_special_method__ = None
    return func

class Cog(metaclass=CogMeta):
    """The base class that all cogs must inherit from.

    A cog is a collection of commands, listeners, and optional state to
    help group commands together. More information on them can be found on
    the :ref:`ext_commands_cogs` page.

    When inheriting from this class, the options shown in :class:`CogMeta`
    are equally valid here.
    """

    def __new__(cls, *args, **kwargs):
        # For issue 426, we need to store a copy of the command objects
        # since we modify them to inject `self` to them.
        # To do this, we need to interfere with the Cog creation process.
        self = super().__new__(cls)
        cmd_attrs = cls.__cog_settings__

        # Either update the command with the cog provided defaults or copy it.
        self.__cog_commands__ = tuple(c._update_copy(cmd_attrs) for c in cls.__cog_commands__)

        lookup = {
            cmd.qualified_name: cmd
            for cmd in self.__cog_commands__
        }

        # Update the Command instances dynamically as well
        for command in self.__cog_commands__:
            setattr(self, command.callback.__name__, command)
            parent = command.parent
            if parent is not None:
                # Get the latest parent reference
                parent = lookup[parent.qualified_name]

                # Update our parent's reference to our self
                removed = parent.remove_command(command.name)
                parent.add_command(command)

        return self

    def get_commands(self):
        r"""
        Returns
        --------
        List[:class:`.Command`]
            A :class:`list` of :class:`.Command`\s that are
            defined inside this cog.

            .. note::

                This does not include subcommands.
        """
        return [c for c in self.__cog_commands__ if c.parent is None]

    @property
    def qualified_name(self):
        """:class:`str`: Returns the cog's specified name, not the class name."""
        return self.__cog_name__

    @property
    def description(self):
        """:class:`str`: Returns the cog's description, typically the cleaned docstring."""
        try:
            return self.__cog_cleaned_doc__
        except AttributeError:
            self.__cog_cleaned_doc__ = cleaned = inspect.getdoc(self)
            return cleaned

    def walk_commands(self):
        """An iterator that recursively walks through this cog's commands and subcommands.

        Yields
        ------
        Union[:class:`.Command`, :class:`.Group`]
            A command or group from the cog.
        """
        from .core import GroupMixin
        for command in self.__cog_commands__:
            if command.parent is None:
                yield command
                if isinstance(command, GroupMixin):
                    yield from command.walk_commands()

    def get_listeners(self):
        """Returns a :class:`list` of (name, function) listener pairs that are defined in this cog.

        Returns
        --------
        List[Tuple[:class:`str`, :ref:`coroutine <coroutine>`]]
            The listeners defined in this cog.
        """
        return [(name, getattr(self, method_name)) for name, method_name in self.__cog_listeners__]

    @classmethod
    def _get_overridden_method(cls, method):
        """Return None if the method is not overridden. Otherwise returns the overridden method."""
        return getattr(method.__func__, '__cog_special_method__', method)

    @classmethod
    def listener(cls, name=None):
        """A decorator that marks a function as a listener.

        This is the cog equivalent of :meth:`.Bot.listen`.

        Parameters
        ------------
        name: :class:`str`
            The name of the event being listened to. If not provided, it
            defaults to the function's name.

        Raises
        --------
        TypeError
            The function is not a coroutine function or a string was not passed as
            the name.
        """

        if name is not None and not isinstance(name, str):
            raise TypeError('Cog.listener expected str but received {0.__class__.__name__!r} instead.'.format(name))

        def decorator(func):
            actual = func
            if isinstance(actual, staticmethod):
                actual = actual.__func__
            if not inspect.iscoroutinefunction(actual):
                raise TypeError('Listener function must be a coroutine function.')
            actual.__cog_listener__ = True
            to_assign = name or actual.__name__
            try:
                actual.__cog_listener_names__.append(to_assign)
            except AttributeError:
                actual.__cog_listener_names__ = [to_assign]
            # we have to return `func` instead of `actual` because
            # we need the type to be `staticmethod` for the metaclass
            # to pick it up but the metaclass unfurls the function and
            # thus the assignments need to be on the actual function
            return func
        return decorator

    @_cog_special_method
    def cog_unload(self):
        """A special method that is called when the cog gets removed.

        This function **cannot** be a coroutine. It must be a regular
        function.

        Subclasses must replace this if they want special unloading behaviour.
        """
        pass

    @_cog_special_method
    def bot_check_once(self, ctx):
        """A special method that registers as a :meth:`.Bot.check_once`
        check.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    def bot_check(self, ctx):
        """A special method that registers as a :meth:`.Bot.check`
        check.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    def cog_check(self, ctx):
        """A special method that registers as a :func:`commands.check`
        for every command and subcommand in this cog.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    async def cog_command_error(self, ctx, error):
        """A special method that is called whenever an error
        is dispatched inside this cog.

        This is similar to :func:`.on_command_error` except only applying
        to the commands inside this cog.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context where the error happened.
        error: :class:`CommandError`
            The error that happened.
        """
        pass

    @_cog_special_method
    async def cog_before_invoke(self, ctx):
        """A special method that acts as a cog local pre-invoke hook.

        This is similar to :meth:`.Command.before_invoke`.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context.
        """
        pass

    @_cog_special_method
    async def cog_after_invoke(self, ctx):
        """A special method that acts as a cog local post-invoke hook.

        This is similar to :meth:`.Command.after_invoke`.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context.
        """
        pass

    def _inject(self, bot):
        cls = self.__class__

        # realistically, the only thing that can cause loading errors
        # is essentially just the command loading, which raises if there are
        # duplicates. When this condition is met, we want to undo all what
        # we've added so far for some form of atomic loading.
        for index, command in enumerate(self.__cog_commands__):
            command.cog = self
            if command.parent is None:
                try:
                    bot.add_command(command)
                except Exception as e:
                    # undo our additions
                    for to_undo in self.__cog_commands__[:index]:
                        bot.remove_command(to_undo.name)
                    raise e

        # check if we're overriding the default
        if cls.bot_check is not Cog.bot_check:
            bot.add_check(self.bot_check)

        if cls.bot_check_once is not Cog.bot_check_once:
            bot.add_check(self.bot_check_once, call_once=True)

        # while Bot.add_listener can raise if it's not a coroutine,
        # this precondition is already met by the listener decorator
        # already, thus this should never raise.
        # Outside of, memory errors and the like...
        for name, method_name in self.__cog_listeners__:
            bot.add_listener(getattr(self, method_name), name)

        return self

    def _eject(self, bot):
        cls = self.__class__

        try:
            for command in self.__cog_commands__:
                if command.parent is None:
                    bot.remove_command(command.name)

            for _, method_name in self.__cog_listeners__:
                bot.remove_listener(getattr(self, method_name))

            if cls.bot_check is not Cog.bot_check:
                bot.remove_check(self.bot_check)

            if cls.bot_check_once is not Cog.bot_check_once:
                bot.remove_check(self.bot_check_once, call_once=True)
        finally:
            self.cog_unload()
<commit_msg>[commands] Allow setting description of cogs<commit_after># -*- coding: utf-8 -*-

"""
The MIT License (MIT)

Copyright (c) 2015-2020 Rapptz

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""

import inspect
import copy
from ._types import _BaseCommand

__all__ = (
    'CogMeta',
    'Cog',
)

class CogMeta(type):
    """A metaclass for defining a cog.

    Note that you should probably not use this directly. It is exposed
    purely for documentation purposes along with making custom metaclasses to intermix
    with other metaclasses such as the :class:`abc.ABCMeta` metaclass.

    For example, to create an abstract cog mixin class, the following would be done.

    .. code-block:: python3

        import abc

        class CogABCMeta(commands.CogMeta, abc.ABCMeta):
            pass

        class SomeMixin(metaclass=abc.ABCMeta):
            pass

        class SomeCogMixin(SomeMixin, commands.Cog, metaclass=CogABCMeta):
            pass

    .. note::

        When passing an attribute of a metaclass that is documented below, note
        that you must pass it as a keyword-only argument to the class creation
        like the following example:

        .. code-block:: python3

            class MyCog(commands.Cog, name='My Cog'):
                pass

    Attributes
    -----------
    name: :class:`str`
        The cog name. By default, it is the name of the class with no modification.
    description: :class:`str`
        The cog description. By default, it is the cleaned docstring of the class.

        .. versionadded:: 1.6

    command_attrs: :class:`dict`
        A list of attributes to apply to every command inside this cog. The dictionary
        is passed into the :class:`Command` options at ``__init__``.
        If you specify attributes inside the command attribute in the class, it will
        override the one specified inside this attribute. For example:

        .. code-block:: python3

            class MyCog(commands.Cog, command_attrs=dict(hidden=True)):
                @commands.command()
                async def foo(self, ctx):
                    pass # hidden -> True

                @commands.command(hidden=False)
                async def bar(self, ctx):
                    pass # hidden -> False
    """

    def __new__(cls, *args, **kwargs):
        name, bases, attrs = args
        attrs['__cog_name__'] = kwargs.pop('name', name)
        attrs['__cog_settings__'] = command_attrs = kwargs.pop('command_attrs', {})

        description = kwargs.pop('description', None)
        if description is None:
            description = inspect.cleandoc(attrs.get('__doc__', ''))
        attrs['__cog_description__'] = description

        commands = {}
        listeners = {}
        no_bot_cog = 'Commands or listeners must not start with cog_ or bot_ (in method {0.__name__}.{1})'

        new_cls = super().__new__(cls, name, bases, attrs, **kwargs)
        for base in reversed(new_cls.__mro__):
            for elem, value in base.__dict__.items():
                if elem in commands:
                    del commands[elem]
                if elem in listeners:
                    del listeners[elem]

                is_static_method = isinstance(value, staticmethod)
                if is_static_method:
                    value = value.__func__
                if isinstance(value, _BaseCommand):
                    if is_static_method:
                        raise TypeError('Command in method {0}.{1!r} must not be staticmethod.'.format(base, elem))
                    if elem.startswith(('cog_', 'bot_')):
                        raise TypeError(no_bot_cog.format(base, elem))
                    commands[elem] = value
                elif inspect.iscoroutinefunction(value):
                    try:
                        is_listener = getattr(value, '__cog_listener__')
                    except AttributeError:
                        continue
                    else:
                        if elem.startswith(('cog_', 'bot_')):
                            raise TypeError(no_bot_cog.format(base, elem))
                        listeners[elem] = value

        new_cls.__cog_commands__ = list(commands.values()) # this will be copied in Cog.__new__

        listeners_as_list = []
        for listener in listeners.values():
            for listener_name in listener.__cog_listener_names__:
                # I use __name__ instead of just storing the value so I can inject
                # the self attribute when the time comes to add them to the bot
                listeners_as_list.append((listener_name, listener.__name__))

        new_cls.__cog_listeners__ = listeners_as_list
        return new_cls

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

    @classmethod
    def qualified_name(cls):
        return cls.__cog_name__

def _cog_special_method(func):
    func.__cog_special_method__ = None
    return func

class Cog(metaclass=CogMeta):
    """The base class that all cogs must inherit from.

    A cog is a collection of commands, listeners, and optional state to
    help group commands together. More information on them can be found on
    the :ref:`ext_commands_cogs` page.

    When inheriting from this class, the options shown in :class:`CogMeta`
    are equally valid here.
    """

    def __new__(cls, *args, **kwargs):
        # For issue 426, we need to store a copy of the command objects
        # since we modify them to inject `self` to them.
        # To do this, we need to interfere with the Cog creation process.
        self = super().__new__(cls)
        cmd_attrs = cls.__cog_settings__

        # Either update the command with the cog provided defaults or copy it.
        self.__cog_commands__ = tuple(c._update_copy(cmd_attrs) for c in cls.__cog_commands__)

        lookup = {
            cmd.qualified_name: cmd
            for cmd in self.__cog_commands__
        }

        # Update the Command instances dynamically as well
        for command in self.__cog_commands__:
            setattr(self, command.callback.__name__, command)
            parent = command.parent
            if parent is not None:
                # Get the latest parent reference
                parent = lookup[parent.qualified_name]

                # Update our parent's reference to our self
                removed = parent.remove_command(command.name)
                parent.add_command(command)

        return self

    def get_commands(self):
        r"""
        Returns
        --------
        List[:class:`.Command`]
            A :class:`list` of :class:`.Command`\s that are
            defined inside this cog.

            .. note::

                This does not include subcommands.
        """
        return [c for c in self.__cog_commands__ if c.parent is None]

    @property
    def qualified_name(self):
        """:class:`str`: Returns the cog's specified name, not the class name."""
        return self.__cog_name__

    @property
    def description(self):
        """:class:`str`: Returns the cog's description, typically the cleaned docstring."""
        return self.__cog_description__

    @description.setter
    def description(self, description):
        self.__cog_description__ = description

    def walk_commands(self):
        """An iterator that recursively walks through this cog's commands and subcommands.

        Yields
        ------
        Union[:class:`.Command`, :class:`.Group`]
            A command or group from the cog.
        """
        from .core import GroupMixin
        for command in self.__cog_commands__:
            if command.parent is None:
                yield command
                if isinstance(command, GroupMixin):
                    yield from command.walk_commands()

    def get_listeners(self):
        """Returns a :class:`list` of (name, function) listener pairs that are defined in this cog.

        Returns
        --------
        List[Tuple[:class:`str`, :ref:`coroutine <coroutine>`]]
            The listeners defined in this cog.
        """
        return [(name, getattr(self, method_name)) for name, method_name in self.__cog_listeners__]

    @classmethod
    def _get_overridden_method(cls, method):
        """Return None if the method is not overridden. Otherwise returns the overridden method."""
        return getattr(method.__func__, '__cog_special_method__', method)

    @classmethod
    def listener(cls, name=None):
        """A decorator that marks a function as a listener.

        This is the cog equivalent of :meth:`.Bot.listen`.

        Parameters
        ------------
        name: :class:`str`
            The name of the event being listened to. If not provided, it
            defaults to the function's name.

        Raises
        --------
        TypeError
            The function is not a coroutine function or a string was not passed as
            the name.
        """

        if name is not None and not isinstance(name, str):
            raise TypeError('Cog.listener expected str but received {0.__class__.__name__!r} instead.'.format(name))

        def decorator(func):
            actual = func
            if isinstance(actual, staticmethod):
                actual = actual.__func__
            if not inspect.iscoroutinefunction(actual):
                raise TypeError('Listener function must be a coroutine function.')
            actual.__cog_listener__ = True
            to_assign = name or actual.__name__
            try:
                actual.__cog_listener_names__.append(to_assign)
            except AttributeError:
                actual.__cog_listener_names__ = [to_assign]
            # we have to return `func` instead of `actual` because
            # we need the type to be `staticmethod` for the metaclass
            # to pick it up but the metaclass unfurls the function and
            # thus the assignments need to be on the actual function
            return func
        return decorator

    @_cog_special_method
    def cog_unload(self):
        """A special method that is called when the cog gets removed.

        This function **cannot** be a coroutine. It must be a regular
        function.

        Subclasses must replace this if they want special unloading behaviour.
        """
        pass

    @_cog_special_method
    def bot_check_once(self, ctx):
        """A special method that registers as a :meth:`.Bot.check_once`
        check.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    def bot_check(self, ctx):
        """A special method that registers as a :meth:`.Bot.check`
        check.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    def cog_check(self, ctx):
        """A special method that registers as a :func:`commands.check`
        for every command and subcommand in this cog.

        This function **can** be a coroutine and must take a sole parameter,
        ``ctx``, to represent the :class:`.Context`.
        """
        return True

    @_cog_special_method
    async def cog_command_error(self, ctx, error):
        """A special method that is called whenever an error
        is dispatched inside this cog.

        This is similar to :func:`.on_command_error` except only applying
        to the commands inside this cog.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context where the error happened.
        error: :class:`CommandError`
            The error that happened.
        """
        pass

    @_cog_special_method
    async def cog_before_invoke(self, ctx):
        """A special method that acts as a cog local pre-invoke hook.

        This is similar to :meth:`.Command.before_invoke`.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context.
        """
        pass

    @_cog_special_method
    async def cog_after_invoke(self, ctx):
        """A special method that acts as a cog local post-invoke hook.

        This is similar to :meth:`.Command.after_invoke`.

        This **must** be a coroutine.

        Parameters
        -----------
        ctx: :class:`.Context`
            The invocation context.
        """
        pass

    def _inject(self, bot):
        cls = self.__class__

        # realistically, the only thing that can cause loading errors
        # is essentially just the command loading, which raises if there are
        # duplicates. When this condition is met, we want to undo all what
        # we've added so far for some form of atomic loading.
        for index, command in enumerate(self.__cog_commands__):
            command.cog = self
            if command.parent is None:
                try:
                    bot.add_command(command)
                except Exception as e:
                    # undo our additions
                    for to_undo in self.__cog_commands__[:index]:
                        bot.remove_command(to_undo.name)
                    raise e

        # check if we're overriding the default
        if cls.bot_check is not Cog.bot_check:
            bot.add_check(self.bot_check)

        if cls.bot_check_once is not Cog.bot_check_once:
            bot.add_check(self.bot_check_once, call_once=True)

        # while Bot.add_listener can raise if it's not a coroutine,
        # this precondition is already met by the listener decorator
        # already, thus this should never raise.
        # Outside of, memory errors and the like...
        for name, method_name in self.__cog_listeners__:
            bot.add_listener(getattr(self, method_name), name)

        return self

    def _eject(self, bot):
        cls = self.__class__

        try:
            for command in self.__cog_commands__:
                if command.parent is None:
                    bot.remove_command(command.name)

            for _, method_name in self.__cog_listeners__:
                bot.remove_listener(getattr(self, method_name))

            if cls.bot_check is not Cog.bot_check:
                bot.remove_check(self.bot_check)

            if cls.bot_check_once is not Cog.bot_check_once:
                bot.remove_check(self.bot_check_once, call_once=True)
        finally:
            self.cog_unload()
<|endoftext|>


<commit_before>from glob import glob
from ccdproc import CCDData
import os
import re
import sys
import numpy as np
import pandas
from pipeline.wcs import WCS
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import logging

class GSPWcsCalculator(object):

    def __init__(self, save_pdf_to=None):
        self.log = logging.getLogger(__name__)
        self.path = None
        self.files = None
        self.ccd = None
        self.pixel = None
        self.angstrom = None
        self.spectrum = None
        self.wcs = WCS()
        self.nist = None
        if save_pdf_to is None:
            self.save_pdf_to = os.getcwd()
        else:
            self.save_pdf_to = save_pdf_to

    def __call__(self, ccd, save=False):
        if isinstance(ccd, CCDData):
            self.ccd = ccd
        else:
            print("Processing file: {:s}".format(ccd))
            self.ccd = CCDData.read(ccd, unit='adu')

        self._recover_lines()
        if not self._validate_lines():
            self.log.error("Please check lines for lamp "
                           "{:s}".format(self.ccd.header['GSP_FNAM']))
            return self.ccd
        else:
            pixel = np.asarray(self.pixel, dtype=float)
            angstrom = np.asarray(self.angstrom, dtype=float)
            wcs_model = self.wcs.fit(physical=pixel,
                                     wavelength=angstrom)
            self.ccd = self.wcs.write_gsp_wcs(self.ccd, wcs_model)
            image_name = os.path.basename(self.ccd.header['GSP_FNAM'])
            # new_name = os.path.join(self.path, image_name[24:])
            new_name = image_name[24:]
            # print(new_name, image_name)
            self.ccd.header.set('GSP_FNAM', value=new_name)
            if save:
                self.ccd.write(new_name, overwrite=True)

            self._pdf_generator(model=wcs_model)
            return self.ccd

    def _recover_lines(self):
        self.pixel = []
        self.angstrom = []
        pixel_keys = self.ccd.header['GSP_P*']
        for pixel_key in pixel_keys:
            if re.match(r'GSP_P\d{3}', pixel_key) is not None:
                angstrom_key = re.sub('GSP_P', 'GSP_A', pixel_key)
                assert pixel_key[-3:] == angstrom_key[-3:]
                assert angstrom_key in self.ccd.header
                if int(self.ccd.header[angstrom_key]) != 0:
                    self.pixel.append(float(self.ccd.header[pixel_key]))
                    self.angstrom.append(float(self.ccd.header[angstrom_key]))
                else:
                    print("File: {:s}".format(self.ccd.header['GSP_FNAM']))
                    print("Ignoring keywords: {:s}={:f}, {:s}={:f}".format(pixel_key, self.ccd.header[pixel_key],
                          angstrom_key, self.ccd.header[angstrom_key]))
        # self.pixel = np.asarray(self.pixel, dtype=float)
        # self.angstrom = np.asarray(self.angstrom, dtype=float)

    def _validate_lines(self):
        """Calls all available validation methods

        Returns:
            True if none of the validation fails.
        """
        assert len(self.pixel) == len(self.angstrom)
        if not self._order_validation(self.pixel):
            return False
        if not self._order_validation(self.angstrom):
            return False
        # if not self._validate_line_existence():
        #     return False
        self._validate_line_existence()
        return True

    @staticmethod
    def _order_validation(lines_array):
        """Checks that the array of lines only increases."""
        previous = None
        for line_value in lines_array:
            # print(line_value)
            if previous is not None:
                try:
                    assert line_value > previous
                    previous = line_value
                except AssertionError:
                    print("Error: Line {:f} is not larger "
                          "than {:f}".format(line_value, previous))
                    return False
            else:
                previous = line_value
        return True

    def _load_nist_list(self, **kwargs):
        """Load all csv files from strong lines in nist."""
        nist_path = kwargs.get(
            'path',
            os.path.join(os.path.dirname(sys.modules['pipeline'].__file__),
                         'data/nist_list'))
        assert os.path.isdir(nist_path)
        nist_files = glob(os.path.join(nist_path, "*.txt"))
        for nist_file in nist_files:
            key = os.path.basename(nist_file)[22:-4]
            nist_data = pandas.read_csv(nist_file, names=['intensity',
                                                          'air_wavelength',
                                                          'spectrum',
                                                          'reference'])
            self.nist[key] = nist_data

    def _validate_line_existence(self):
        """Check if a line actually exists in any table of NIST
        Notes:
            It does not actually check NIST, it loads six csv tables
            from NIST's strong lines for the elements used in lamps.
            Ar,  Cu, Fe, He, Hg, Ne. It does not work perfect so far
            so the it is not checking existence actually but if it finds it
            it will get the value at "spectrum" column in NIST tables which
            correspond to the source of the line for instance Ne I, Ar II.
            """

        lamp_elements = []
        lamp_name = self.ccd.header['OBJECT']
        if len(lamp_name) % 2 == 0:
            for element_index in range(0, len(lamp_name), 2):
                element = lamp_name[element_index:element_index + 2]
                lamp_elements.append(element)

        if self.nist is None:
            self.nist = {}
            self._load_nist_list()

        self.spectrum = list(self.angstrom)
        for i in range(len(self.angstrom)):
            for element in lamp_elements:
                line_info = self.nist[element][
                    self.nist[element].air_wavelength == self.angstrom[i]]
                if line_info.empty:
                    # print(self.angstrom[i], 'no-info')
                    self.spectrum[i] = ''
                else:
                    self.spectrum[i] = line_info['spectrum'].to_string(index=False)
                    # print(self.angstrom[i], line_info['spectrum'].to_string(index=False))
            print(self.angstrom[i], self.spectrum[i])

    def _pdf_generator(self, model):
        """Creates a pdf file Using Reference lines."""
        file_name = self.ccd.header['GSP_FNAM']
        pdf_file_name = os.path.join(self.save_pdf_to,
                                     re.sub('.fits', '.pdf', file_name))

        with PdfPages(pdf_file_name) as pdf:
            plt.figure(1, (40, 10))
            plt.plot(model(range(len(self.ccd.data))), self.ccd.data, color='k')
            line_spacer = 0

            top_lim = 1.3 * self.ccd.data.max()
            bottom_lim = self.ccd.data.min() - 0.05 * self.ccd.data.max()
            plt.ylim((bottom_lim, top_lim))
            plt.xlim((model(0), model(len(self.ccd.data))))

            for i in range(len(self.angstrom)):
                line_str = "{:.4f} {:s}".format(self.angstrom[i], str(self.spectrum[i]))
                try:
                    if self.angstrom[i+1] - self.angstrom[i] < 25:
                        line_spacer = (25 - (self.angstrom[i+1] - self.angstrom[i]))
                        y_spacer = 0.1 * self.ccd.data.max()
                    elif line_spacer > 0:
                        line_spacer *= -1
                    else:
                        line_spacer = 0
                        y_spacer = 0
                    # print(line_spacer)
                except IndexError:
                    line_spacer = 0
                    y_spacer = 0
                x_pos = self.angstrom[i]
                y_pos = np.max((self.ccd.data[int(np.floor(self.pixel[i]))],
                                self.ccd.data[int(np.ceil(self.pixel[i]))]))
                y_offset = 0.05 * self.ccd.data.max()
                # plt.plot(x_pos, y_pos, x_pos - line_spacer, y_pos + y_offset, color='r')
                # y_min = y_pos / self.ccd.data.max()
                # y_max = (y_pos + y_offset) / self.ccd.data.max()
                # plt.axvline(x=x_pos, ymin=y_min, ymax=y_max, color='r')
                plt.text(x_pos, y_pos + y_offset, line_str, rotation=90,
                         verticalalignment='bottom',
                         horizontalalignment='center')
                # plt.axvline(aline, color='r', alpha=0.5)

            title = "{:s} - {:s} - {:s}".format(self.ccd.header['OBJECT'],
                                                self.ccd.header['WAVMODE'],
                                                self.ccd.header['SLIT'])
            plt.title(title)
            plt.xlabel('Wavelength (Angstrom)')
            plt.ylabel('Intensity (ADU)')
            plt.tight_layout()
            pdf.savefig()
            # pdf.close()
            # plt.show()
            plt.close()

if __name__ == '__main__':
    calculator = GSPWcsCalculator()
    if len(sys.argv) == 2:
        _file = sys.argv[1]
        if not os.path.isabs(_file):
            _file = os.path.join(os.getcwd(), _file)
        ccd = CCDData.read(_file, unit='adu')

        calculator(ccd=ccd, save=True)
    else:
        sys.exit("You must provide a file name\n"
                 "{:s} file-name.fits".format(sys.argv[0]))
<commit_msg>Corrected convertion to integer<commit_after>from glob import glob
from ccdproc import CCDData
import os
import re
import sys
import numpy as np
import pandas
from pipeline.wcs import WCS
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import logging

class GSPWcsCalculator(object):

    def __init__(self, save_pdf_to=None):
        self.log = logging.getLogger(__name__)
        self.path = None
        self.files = None
        self.ccd = None
        self.pixel = None
        self.angstrom = None
        self.spectrum = None
        self.wcs = WCS()
        self.nist = None
        if save_pdf_to is None:
            self.save_pdf_to = os.getcwd()
        else:
            self.save_pdf_to = save_pdf_to

    def __call__(self, ccd, save=False):
        if isinstance(ccd, CCDData):
            self.ccd = ccd
        else:
            print("Processing file: {:s}".format(ccd))
            self.ccd = CCDData.read(ccd, unit='adu')

        self._recover_lines()
        if not self._validate_lines():
            self.log.error("Please check lines for lamp "
                           "{:s}".format(self.ccd.header['GSP_FNAM']))
            return self.ccd
        else:
            pixel = np.asarray(self.pixel, dtype=float)
            angstrom = np.asarray(self.angstrom, dtype=float)
            wcs_model = self.wcs.fit(physical=pixel,
                                     wavelength=angstrom)
            self.ccd = self.wcs.write_gsp_wcs(self.ccd, wcs_model)
            image_name = os.path.basename(self.ccd.header['GSP_FNAM'])
            # new_name = os.path.join(self.path, image_name[24:])
            new_name = image_name[24:]
            # print(new_name, image_name)
            self.ccd.header.set('GSP_FNAM', value=new_name)
            if save:
                self.ccd.write(new_name, overwrite=True)

            self._pdf_generator(model=wcs_model)
            return self.ccd

    def _recover_lines(self):
        self.pixel = []
        self.angstrom = []
        pixel_keys = self.ccd.header['GSP_P*']
        for pixel_key in pixel_keys:
            if re.match(r'GSP_P\d{3}', pixel_key) is not None:
                angstrom_key = re.sub('GSP_P', 'GSP_A', pixel_key)
                assert pixel_key[-3:] == angstrom_key[-3:]
                assert angstrom_key in self.ccd.header
                if int(float(self.ccd.header[angstrom_key])) != 0:
                    self.pixel.append(float(self.ccd.header[pixel_key]))
                    self.angstrom.append(float(self.ccd.header[angstrom_key]))
                else:
                    print("File: {:s}".format(self.ccd.header['GSP_FNAM']))
                    print("Ignoring keywords: {:s}={:f}, {:s}={:f}".format(pixel_key, self.ccd.header[pixel_key],
                          angstrom_key, float(self.ccd.header[angstrom_key])))
        # self.pixel = np.asarray(self.pixel, dtype=float)
        # self.angstrom = np.asarray(self.angstrom, dtype=float)

    def _validate_lines(self):
        """Calls all available validation methods

        Returns:
            True if none of the validation fails.
        """
        assert len(self.pixel) == len(self.angstrom)
        if not self._order_validation(self.pixel):
            return False
        if not self._order_validation(self.angstrom):
            return False
        # if not self._validate_line_existence():
        #     return False
        self._validate_line_existence()
        return True

    @staticmethod
    def _order_validation(lines_array):
        """Checks that the array of lines only increases."""
        previous = None
        for line_value in lines_array:
            # print(line_value)
            if previous is not None:
                try:
                    assert line_value > previous
                    previous = line_value
                except AssertionError:
                    print("Error: Line {:f} is not larger "
                          "than {:f}".format(line_value, previous))
                    return False
            else:
                previous = line_value
        return True

    def _load_nist_list(self, **kwargs):
        """Load all csv files from strong lines in nist."""
        nist_path = kwargs.get(
            'path',
            os.path.join(os.path.dirname(sys.modules['pipeline'].__file__),
                         'data/nist_list'))
        assert os.path.isdir(nist_path)
        nist_files = glob(os.path.join(nist_path, "*.txt"))
        for nist_file in nist_files:
            key = os.path.basename(nist_file)[22:-4]
            nist_data = pandas.read_csv(nist_file, names=['intensity',
                                                          'air_wavelength',
                                                          'spectrum',
                                                          'reference'])
            self.nist[key] = nist_data

    def _validate_line_existence(self):
        """Check if a line actually exists in any table of NIST
        Notes:
            It does not actually check NIST, it loads six csv tables
            from NIST's strong lines for the elements used in lamps.
            Ar,  Cu, Fe, He, Hg, Ne. It does not work perfect so far
            so the it is not checking existence actually but if it finds it
            it will get the value at "spectrum" column in NIST tables which
            correspond to the source of the line for instance Ne I, Ar II.
            """

        lamp_elements = []
        lamp_name = self.ccd.header['OBJECT']
        if len(lamp_name) % 2 == 0:
            for element_index in range(0, len(lamp_name), 2):
                element = lamp_name[element_index:element_index + 2]
                lamp_elements.append(element)

        if self.nist is None:
            self.nist = {}
            self._load_nist_list()

        self.spectrum = list(self.angstrom)
        for i in range(len(self.angstrom)):
            for element in lamp_elements:
                line_info = self.nist[element][
                    self.nist[element].air_wavelength == self.angstrom[i]]
                if line_info.empty:
                    # print(self.angstrom[i], 'no-info')
                    self.spectrum[i] = ''
                else:
                    self.spectrum[i] = line_info['spectrum'].to_string(index=False)
                    # print(self.angstrom[i], line_info['spectrum'].to_string(index=False))
            print(self.angstrom[i], self.spectrum[i])

    def _pdf_generator(self, model):
        """Creates a pdf file Using Reference lines."""
        file_name = self.ccd.header['GSP_FNAM']
        pdf_file_name = os.path.join(self.save_pdf_to,
                                     re.sub('.fits', '.pdf', file_name))

        with PdfPages(pdf_file_name) as pdf:
            plt.figure(1, (40, 10))
            plt.plot(model(range(len(self.ccd.data))), self.ccd.data, color='k')
            line_spacer = 0

            top_lim = 1.3 * self.ccd.data.max()
            bottom_lim = self.ccd.data.min() - 0.05 * self.ccd.data.max()
            plt.ylim((bottom_lim, top_lim))
            plt.xlim((model(0), model(len(self.ccd.data))))

            for i in range(len(self.angstrom)):
                line_str = "{:.4f} {:s}".format(self.angstrom[i], str(self.spectrum[i]))
                try:
                    if self.angstrom[i+1] - self.angstrom[i] < 25:
                        line_spacer = (25 - (self.angstrom[i+1] - self.angstrom[i]))
                        y_spacer = 0.1 * self.ccd.data.max()
                    elif line_spacer > 0:
                        line_spacer *= -1
                    else:
                        line_spacer = 0
                        y_spacer = 0
                    # print(line_spacer)
                except IndexError:
                    line_spacer = 0
                    y_spacer = 0
                x_pos = self.angstrom[i]
                y_pos = np.max((self.ccd.data[int(np.floor(self.pixel[i]))],
                                self.ccd.data[int(np.ceil(self.pixel[i]))]))
                y_offset = 0.05 * self.ccd.data.max()
                # plt.plot(x_pos, y_pos, x_pos - line_spacer, y_pos + y_offset, color='r')
                # y_min = y_pos / self.ccd.data.max()
                # y_max = (y_pos + y_offset) / self.ccd.data.max()
                # plt.axvline(x=x_pos, ymin=y_min, ymax=y_max, color='r')
                plt.text(x_pos, y_pos + y_offset, line_str, rotation=90,
                         verticalalignment='bottom',
                         horizontalalignment='center')
                # plt.axvline(aline, color='r', alpha=0.5)

            title = "{:s} - {:s} - {:s}".format(self.ccd.header['OBJECT'],
                                                self.ccd.header['WAVMODE'],
                                                self.ccd.header['SLIT'])
            plt.title(title)
            plt.xlabel('Wavelength (Angstrom)')
            plt.ylabel('Intensity (ADU)')
            plt.tight_layout()
            pdf.savefig()
            # pdf.close()
            # plt.show()
            plt.close()

if __name__ == '__main__':
    calculator = GSPWcsCalculator()
    if len(sys.argv) == 2:
        _file = sys.argv[1]
        if not os.path.isabs(_file):
            _file = os.path.join(os.getcwd(), _file)
        ccd = CCDData.read(_file, unit='adu')

        calculator(ccd=ccd, save=True)
    else:
        sys.exit("You must provide a file name\n"
                 "{:s} file-name.fits".format(sys.argv[0]))
<|endoftext|>


<commit_before>#!/usr/bin/env python
import os
import sys
import simdna
import simdna.util as util
import simdna.synthetic as synthetic
import simdna.pwm as pwm

def variableSpacingGrammar(options):
    pc = 0.001
    pathToMotifs = options.pathToMotifs
    loadedMotifs = synthetic.LoadedEncodeMotifs(pathToMotifs, pseudocountProb=pc)
    motifName1 = options.motifName1
    motifName2 = options.motifName2
    seqLength = options.seqLength
    numSeq = options.numSeq
    outputFileName = ("variableSpacingGrammarSimulation_motif1-"
                      +motifName1+"_motif2-"+motifName2
                      +"_seqLength"+str(seqLength)+"_numSeq"
                      +str(numSeq)+".simdata")

    kwargs={'loadedMotifs':loadedMotifs}
    theClass=synthetic.PwmSamplerFromLoadedMotifs
    motif1Generator=theClass(motifName=motifName1,**kwargs)
    motif2Generator=theClass(motifName=motifName2,**kwargs)
    motif1Embedder=synthetic.SubstringEmbedder(substringGenerator=motif1Generator)
    motif2Embedder=synthetic.SubstringEmbedder(substringGenerator=motif2Generator)

    embedders = []
    namePrefix="synthPos"
    separationGenerator=synthetic.MinMaxWrapper(
        synthetic.PoissonQuantityGenerator(options.meanSpacing),
        theMin=options.minSpacing,
        theMax=options.maxSpacing)
    embedders.append(synthetic.EmbeddableEmbedder(
                        embeddableGenerator=synthetic.PairEmbeddableGenerator(
                            embeddableGenerator1=motif1Generator
                            ,embeddableGenerator2=motif2Generator
                            ,separationGenerator=separationGenerator
                        )
                    ))

    embedInBackground = synthetic.EmbedInABackground(
        backgroundGenerator=synthetic.ZeroOrderBackgroundGenerator(seqLength)
        , embedders=embedders
        , namePrefix=namePrefix
    )

    sequenceSet = synthetic.GenerateSequenceNTimes(embedInBackground, numSeq)
    synthetic.printSequences(outputFileName, sequenceSet,
                             includeFasta=True, includeEmbeddings=True)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--pathToMotifs",
        default=simdna.ENCODE_MOTIFS_PATH)
    parser.add_argument("--motifName1", required=True)
    parser.add_argument("--motifName2", required=True)
    parser.add_argument("--seqLength", type=int, required=True)
    parser.add_argument("--numSeq", type=int, required=True)
    parser.add_argument("--minSpacing", type=int, required=True)
    parser.add_argument("--meanSpacing", type=float, required=True)
    parser.add_argument("--maxSpacing", type=int, required=True)

    options = parser.parse_args()
    variableSpacingGrammar(options)
<commit_msg>added prefix<commit_after>#!/usr/bin/env python
import os
import sys
import simdna
import simdna.util as util
import simdna.synthetic as synthetic
import simdna.pwm as pwm

def variableSpacingGrammar(options):
    pc = 0.001
    pathToMotifs = options.pathToMotifs
    loadedMotifs = synthetic.LoadedEncodeMotifs(pathToMotifs, pseudocountProb=pc)
    motifName1 = options.motifName1
    motifName2 = options.motifName2
    seqLength = options.seqLength
    numSeq = options.numSeq
    outputFileName = ("variableSpacingGrammarSimulation_"
                      +"prefix-"+options.prefix
                      +"_motif1-"+motifName1+"_motif2-"+motifName2
                      +"_seqLength"+str(seqLength)+"_numSeq"
                      +str(numSeq)+".simdata")

    kwargs={'loadedMotifs':loadedMotifs}
    theClass=synthetic.PwmSamplerFromLoadedMotifs
    motif1Generator=theClass(motifName=motifName1,**kwargs)
    motif2Generator=theClass(motifName=motifName2,**kwargs)
    motif1Embedder=synthetic.SubstringEmbedder(substringGenerator=motif1Generator)
    motif2Embedder=synthetic.SubstringEmbedder(substringGenerator=motif2Generator)

    embedders = []
    separationGenerator=synthetic.MinMaxWrapper(
        synthetic.PoissonQuantityGenerator(options.meanSpacing),
        theMin=options.minSpacing,
        theMax=options.maxSpacing)
    embedders.append(synthetic.EmbeddableEmbedder(
                        embeddableGenerator=synthetic.PairEmbeddableGenerator(
                            embeddableGenerator1=motif1Generator
                            ,embeddableGenerator2=motif2Generator
                            ,separationGenerator=separationGenerator
                        )
                    ))

    embedInBackground = synthetic.EmbedInABackground(
        backgroundGenerator=synthetic.ZeroOrderBackgroundGenerator(seqLength)
        , embedders=embedders
    )

    sequenceSet = synthetic.GenerateSequenceNTimes(embedInBackground, numSeq)
    synthetic.printSequences(outputFileName, sequenceSet,
                             includeFasta=True, includeEmbeddings=True,
                             prefix=options.prefix)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--prefix",
        default="")
    parser.add_argument("--pathToMotifs",
        default=simdna.ENCODE_MOTIFS_PATH)
    parser.add_argument("--motifName1", required=True)
    parser.add_argument("--motifName2", required=True)
    parser.add_argument("--seqLength", type=int, required=True)
    parser.add_argument("--numSeq", type=int, required=True)
    parser.add_argument("--minSpacing", type=int, required=True)
    parser.add_argument("--meanSpacing", type=float, required=True)
    parser.add_argument("--maxSpacing", type=int, required=True)

    options = parser.parse_args()
    variableSpacingGrammar(options)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, unicode_literals

import logging
import os
from io import BytesIO

from django.core.files.base import ContentFile
from django.core.files.storage import DefaultStorage
from django.db.models import signals
from django.db.models.fields.files import (ImageField, ImageFieldFile,
                                           ImageFileDescriptor)
from PIL import Image, ImageOps

from .validators import MinSizeValidator

logger = logging.getLogger()

class StdImageFileDescriptor(ImageFileDescriptor):

    """The variation property of the field is accessible in instance cases."""

    def __set__(self, instance, value):
        super(StdImageFileDescriptor, self).__set__(instance, value)
        self.field.set_variations(instance)

class StdImageFieldFile(ImageFieldFile):

    """Like ImageFieldFile but handles variations."""

    def save(self, name, content, save=True):
        super(StdImageFieldFile, self).save(name, content, save)
        render_variations = self.field.render_variations
        if callable(render_variations):
            render_variations = render_variations(
                file_name=self.name,
                variations=self.field.variations,
                storage=self.storage,
            )
        if not isinstance(render_variations, bool):
            msg = (
                '"render_variations" callable expects a boolean return value,'
                ' but got %s'
                ) % type(render_variations)
            raise TypeError(msg)
        if render_variations:
            self.render_variations()

    @staticmethod
    def is_smaller(img, variation):
        return img.size[0] > variation['width'] \
            or img.size[1] > variation['height']

    def render_variations(self, replace=False):
        """Render all image variations and saves them to the storage."""
        for _, variation in self.field.variations.items():
            self.render_variation(self.name, variation, replace, self.storage)

    @classmethod
    def render_variation(cls, file_name, variation, replace=False,
                         storage=DefaultStorage()):
        """Render an image variation and saves it to the storage."""
        variation_name = cls.get_variation_name(file_name, variation['name'])
        if storage.exists(variation_name):
            if replace:
                storage.delete(variation_name)
                logger.info('File "{}" already exists and has been replaced.')
            else:
                logger.info('File "{}" already exists.')
                return variation_name

        resample = variation['resample']

        with storage.open(file_name) as f:
            with Image.open(f) as img:
                file_format = img.format

                if cls.is_smaller(img, variation):
                    factor = 1
                    while img.size[0] / factor \
                            > 2 * variation['width'] \
                            and img.size[1] * 2 / factor \
                            > 2 * variation['height']:
                        factor *= 2
                    if factor > 1:
                        img.thumbnail(
                            (int(img.size[0] / factor),
                             int(img.size[1] / factor)),
                            resample=resample
                        )

                    size = variation['width'], variation['height']
                    size = tuple(int(i) if i != float('inf') else i
                                 for i in size)
                    if variation['crop']:
                        img = ImageOps.fit(
                            img,
                            size,
                            method=resample
                        )
                    else:
                        img.thumbnail(
                            size,
                            resample=resample
                        )

                with BytesIO() as file_buffer:
                    img.save(file_buffer, file_format)
                    f = ContentFile(file_buffer.getvalue())
                    storage.save(variation_name, f)
        return variation_name

    @classmethod
    def get_variation_name(cls, file_name, variation_name):
        """Return the variation file name based on the variation."""
        path, ext = os.path.splitext(file_name)
        path, file_name = os.path.split(path)
        file_name = '{file_name}.{variation_name}{extension}'.format(**{
            'file_name': file_name,
            'variation_name': variation_name,
            'extension': ext,
        })
        return os.path.join(path, file_name)

    def delete(self, save=True):
        self.delete_variations()
        super(StdImageFieldFile, self).delete(save)

    def delete_variations(self):
        for variation in self.field.variations:
            variation_name = self.get_variation_name(self.name, variation)
            self.storage.delete(variation_name)

class StdImageField(ImageField):

    """
    Django ImageField that is able to create different size variations.

    Extra features are:
        - Django-Storages compatible (S3)
        - Python 2, 3 and PyPy support
        - Django 1.5 and later support
        - Resize images to different sizes
        - Access thumbnails on model level, no template tags required
        - Preserves original image
        - Asynchronous rendering (Celery & Co)
        - Multi threading and processing for optimum performance
        - Restrict accepted image dimensions
        - Rename files to a standardized name (using a callable upload_to)

    :param variations: size variations of the image
    """

    descriptor_class = StdImageFileDescriptor
    attr_class = StdImageFieldFile
    def_variation = {
        'width': float('inf'),
        'height': float('inf'),
        'crop': False,
        'resample': Image.ANTIALIAS
    }

    def __init__(self, verbose_name=None, name=None, variations=None,
                 render_variations=True, force_min_size=False,
                 *args, **kwargs):
        """
        Standardized ImageField for Django.

        Usage: StdImageField(upload_to='PATH',
         variations={'thumbnail': {"width", "height", "crop", "resample"}})
        :param variations: size variations of the image
        :rtype variations: StdImageField
        :param render_variations: boolean or callable that returns a boolean.
         The callable gets passed the app_name, model, field_name and pk.
         Default: True
        :rtype render_variations: bool, callable
        """
        if not variations:
            variations = {}
        if not isinstance(variations, dict):
            msg = ('"variations" expects a dict,'
                   ' but got %s') % type(variations)
            raise TypeError(msg)
        if not (isinstance(render_variations, bool) or
                callable(render_variations)):
            msg = ('"render_variations" excepts a boolean or callable,'
                   ' but got %s') % type(render_variations)
            raise TypeError(msg)

        self._variations = variations
        self.force_min_size = force_min_size
        self.render_variations = render_variations
        self.variations = {}

        for nm, prm in list(variations.items()):
            self.add_variation(nm, prm)

        if self.variations and self.force_min_size:
            self.min_size = (
                max(self.variations.values(),
                    key=lambda x: x["width"])["width"],
                max(self.variations.values(),
                    key=lambda x: x["height"])["height"]
            )

        super(StdImageField, self).__init__(verbose_name, name,
                                            *args, **kwargs)

    def add_variation(self, name, params):
        variation = self.def_variation.copy()
        if isinstance(params, (list, tuple)):
            variation.update(dict(zip(("width", "height", "crop"), params)))
        else:
            variation.update(params)
        variation["name"] = name
        self.variations[name] = variation

    def set_variations(self, instance=None, **kwargs):
        """
        Create a "variation" object as attribute of the ImageField instance.

        Variation attribute will be of the same class as the original image, so
        "path", "url"... properties can be used

        :param instance: FileField
        """
        if getattr(instance, self.name):
            field = getattr(instance, self.name)
            if field._committed:
                for name, variation in list(self.variations.items()):
                    variation_name = self.attr_class.get_variation_name(
                        field.name,
                        variation['name']
                    )
                    variation_field = ImageFieldFile(instance,
                                                     self,
                                                     variation_name)
                    setattr(field, name, variation_field)

    def contribute_to_class(self, cls, name):
        """Generating all operations on specified signals."""
        super(StdImageField, self).contribute_to_class(cls, name)
        signals.post_init.connect(self.set_variations, sender=cls)

    def validate(self, value, model_instance):
        super(StdImageField, self).validate(value, model_instance)
        if self.force_min_size:
            MinSizeValidator(self.min_size[0], self.min_size[1])(value)
<commit_msg>Hack crop options<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, unicode_literals

import logging
import os
from io import BytesIO

from django.core.files.base import ContentFile
from django.core.files.storage import DefaultStorage
from django.db.models import signals
from django.db.models.fields.files import (ImageField, ImageFieldFile,
                                           ImageFileDescriptor)
from PIL import Image, ImageOps

from .validators import MinSizeValidator

logger = logging.getLogger()

class StdImageFileDescriptor(ImageFileDescriptor):

    """The variation property of the field is accessible in instance cases."""

    def __set__(self, instance, value):
        super(StdImageFileDescriptor, self).__set__(instance, value)
        self.field.set_variations(instance)

class StdImageFieldFile(ImageFieldFile):

    """Like ImageFieldFile but handles variations."""

    def save(self, name, content, save=True):
        super(StdImageFieldFile, self).save(name, content, save)
        render_variations = self.field.render_variations
        if callable(render_variations):
            render_variations = render_variations(
                file_name=self.name,
                variations=self.field.variations,
                storage=self.storage,
            )
        if not isinstance(render_variations, bool):
            msg = (
                '"render_variations" callable expects a boolean return value,'
                ' but got %s'
            ) % type(render_variations)
            raise TypeError(msg)
        if render_variations:
            self.render_variations()

    @staticmethod
    def is_smaller(img, variation):
        return img.size[0] > variation['width'] \
            or img.size[1] > variation['height']

    def render_variations(self, replace=False):
        """Render all image variations and saves them to the storage."""
        for _, variation in self.field.variations.items():
            self.render_variation(self.name, variation, replace, self.storage)

    @classmethod
    def render_variation(cls, file_name, variation, replace=False,
                         storage=DefaultStorage()):
        """Render an image variation and saves it to the storage."""
        variation_name = cls.get_variation_name(file_name, variation['name'])
        if storage.exists(variation_name):
            if replace:
                storage.delete(variation_name)
                logger.info('File "{}" already exists and has been replaced.')
            else:
                logger.info('File "{}" already exists.')
                return variation_name

        resample = variation['resample']

        with storage.open(file_name) as f:
            with Image.open(f) as img:
                file_format = img.format

                if cls.is_smaller(img, variation):
                    factor = 1
                    while img.size[0] / factor \
                            > 2 * variation['width'] \
                            and img.size[1] * 2 / factor \
                            > 2 * variation['height']:
                        factor *= 2
                    if factor > 1:
                        img.thumbnail(
                            (int(img.size[0] / factor),
                             int(img.size[1] / factor)),
                            resample=resample
                        )

                    size = variation['width'], variation['height']
                    size = tuple(int(i) if i != float('inf') else i
                                 for i in size)
                    if variation['crop']:
                        img = ImageOps.fit(
                            img,
                            size,
                            method=resample
                        )
                    else:
                        img.thumbnail(
                            size,
                            resample=resample
                        )
                '''
                Kowit Hacking to make crop options fit image anyway
                '''
                size = variation['width'], variation['height']
                size = tuple(int(i) if i != float('inf') else i
                             for i in size)
                if variation['crop']:
                    img = ImageOps.fit(
                        img,
                        size,
                        method=resample
                    )
                '''
                TODO : Hack in in correct way.
                '''
                with BytesIO() as file_buffer:
                    img.save(file_buffer, file_format)
                    f = ContentFile(file_buffer.getvalue())
                    storage.save(variation_name, f)
        return variation_name

    @classmethod
    def get_variation_name(cls, file_name, variation_name):
        """Return the variation file name based on the variation."""
        path, ext = os.path.splitext(file_name)
        path, file_name = os.path.split(path)
        file_name = '{file_name}.{variation_name}{extension}'.format(**{
            'file_name': file_name,
            'variation_name': variation_name,
            'extension': ext,
        })
        return os.path.join(path, file_name)

    def delete(self, save=True):
        self.delete_variations()
        super(StdImageFieldFile, self).delete(save)

    def delete_variations(self):
        for variation in self.field.variations:
            variation_name = self.get_variation_name(self.name, variation)
            self.storage.delete(variation_name)

class StdImageField(ImageField):

    """
    Django ImageField that is able to create different size variations.

    Extra features are:
        - Django-Storages compatible (S3)
        - Python 2, 3 and PyPy support
        - Django 1.5 and later support
        - Resize images to different sizes
        - Access thumbnails on model level, no template tags required
        - Preserves original image
        - Asynchronous rendering (Celery & Co)
        - Multi threading and processing for optimum performance
        - Restrict accepted image dimensions
        - Rename files to a standardized name (using a callable upload_to)

    :param variations: size variations of the image
    """

    descriptor_class = StdImageFileDescriptor
    attr_class = StdImageFieldFile
    def_variation = {
        'width': float('inf'),
        'height': float('inf'),
        'crop': False,
        'resample': Image.ANTIALIAS
    }

    def __init__(self, verbose_name=None, name=None, variations=None,
                 render_variations=True, force_min_size=False,
                 *args, **kwargs):
        """
        Standardized ImageField for Django.

        Usage: StdImageField(upload_to='PATH',
         variations={'thumbnail': {"width", "height", "crop", "resample"}})
        :param variations: size variations of the image
        :rtype variations: StdImageField
        :param render_variations: boolean or callable that returns a boolean.
         The callable gets passed the app_name, model, field_name and pk.
         Default: True
        :rtype render_variations: bool, callable
        """
        if not variations:
            variations = {}
        if not isinstance(variations, dict):
            msg = ('"variations" expects a dict,'
                   ' but got %s') % type(variations)
            raise TypeError(msg)
        if not (isinstance(render_variations, bool) or
                callable(render_variations)):
            msg = ('"render_variations" excepts a boolean or callable,'
                   ' but got %s') % type(render_variations)
            raise TypeError(msg)

        self._variations = variations
        self.force_min_size = force_min_size
        self.render_variations = render_variations
        self.variations = {}

        for nm, prm in list(variations.items()):
            self.add_variation(nm, prm)

        if self.variations and self.force_min_size:
            self.min_size = (
                max(self.variations.values(),
                    key=lambda x: x["width"])["width"],
                max(self.variations.values(),
                    key=lambda x: x["height"])["height"]
            )

        super(StdImageField, self).__init__(verbose_name, name,
                                            *args, **kwargs)

    def add_variation(self, name, params):
        variation = self.def_variation.copy()
        if isinstance(params, (list, tuple)):
            variation.update(dict(zip(("width", "height", "crop"), params)))
        else:
            variation.update(params)
        variation["name"] = name
        self.variations[name] = variation

    def set_variations(self, instance=None, **kwargs):
        """
        Create a "variation" object as attribute of the ImageField instance.

        Variation attribute will be of the same class as the original image, so
        "path", "url"... properties can be used

        :param instance: FileField
        """
        if getattr(instance, self.name):
            field = getattr(instance, self.name)
            if field._committed:
                for name, variation in list(self.variations.items()):
                    variation_name = self.attr_class.get_variation_name(
                        field.name,
                        variation['name']
                    )
                    variation_field = ImageFieldFile(instance,
                                                     self,
                                                     variation_name)
                    setattr(field, name, variation_field)

    def contribute_to_class(self, cls, name):
        """Generating all operations on specified signals."""
        super(StdImageField, self).contribute_to_class(cls, name)
        signals.post_init.connect(self.set_variations, sender=cls)

    def validate(self, value, model_instance):
        super(StdImageField, self).validate(value, model_instance)
        if self.force_min_size:
            MinSizeValidator(self.min_size[0], self.min_size[1])(value)
<|endoftext|>


<commit_before># Copyright (c) OpenMMLab. All rights reserved.
import torch

from ..builder import BBOX_SAMPLERS
from .base_sampler import BaseSampler
from .sampling_result import SamplingResult

@BBOX_SAMPLERS.register_module()
class PseudoSampler(BaseSampler):
    """A pseudo sampler that does not do sampling actually."""

    def __init__(self, **kwargs):
        pass

    def _sample_pos(self, **kwargs):
        """Sample positive samples."""
        raise NotImplementedError

    def _sample_neg(self, **kwargs):
        """Sample negative samples."""
        raise NotImplementedError

    def sample(self, assign_result, bboxes, gt_bboxes, **kwargs):
        """Directly returns the positive and negative indices  of samples.

        Args:
            assign_result (:obj:`AssignResult`): Assigned results
            bboxes (torch.Tensor): Bounding boxes
            gt_bboxes (torch.Tensor): Ground truth boxes

        Returns:
            :obj:`SamplingResult`: sampler results
        """
        pos_inds = torch.nonzero(
            assign_result.gt_inds > 0, as_tuple=False).squeeze(-1).unique()
        neg_inds = torch.nonzero(
            assign_result.gt_inds == 0, as_tuple=False).squeeze(-1).unique()
        gt_flags = bboxes.new_zeros(bboxes.shape[0], dtype=torch.uint8)
        sampling_result = SamplingResult(pos_inds, neg_inds, bboxes, gt_bboxes,
                                         assign_result, gt_flags)
        return sampling_result
<commit_msg>fix pseudosampler (#6622)<commit_after># Copyright (c) OpenMMLab. All rights reserved.
import torch

from ..builder import BBOX_SAMPLERS
from .base_sampler import BaseSampler
from .sampling_result import SamplingResult

@BBOX_SAMPLERS.register_module()
class PseudoSampler(BaseSampler):
    """A pseudo sampler that does not do sampling actually."""

    def __init__(self, **kwargs):
        pass

    def _sample_pos(self, **kwargs):
        """Sample positive samples."""
        raise NotImplementedError

    def _sample_neg(self, **kwargs):
        """Sample negative samples."""
        raise NotImplementedError

    def sample(self, assign_result, bboxes, gt_bboxes, *args, **kwargs):
        """Directly returns the positive and negative indices  of samples.

        Args:
            assign_result (:obj:`AssignResult`): Assigned results
            bboxes (torch.Tensor): Bounding boxes
            gt_bboxes (torch.Tensor): Ground truth boxes

        Returns:
            :obj:`SamplingResult`: sampler results
        """
        pos_inds = torch.nonzero(
            assign_result.gt_inds > 0, as_tuple=False).squeeze(-1).unique()
        neg_inds = torch.nonzero(
            assign_result.gt_inds == 0, as_tuple=False).squeeze(-1).unique()
        gt_flags = bboxes.new_zeros(bboxes.shape[0], dtype=torch.uint8)
        sampling_result = SamplingResult(pos_inds, neg_inds, bboxes, gt_bboxes,
                                         assign_result, gt_flags)
        return sampling_result
<|endoftext|>


<commit_before>import os
import requests
import json

BASE_URL = "https://api.majestic.com/api/json"
BASE_PARAMS = {'app_api_key': os.environ.get('THEAPIKEY')}

def get(cmd, params):
    querydict = {'cmd': cmd}
    querydict.update(BASE_PARAMS)
    querydict.update(params)
    response = requests.get(BASE_URL, params=querydict)
    return json.loads(response.text)

if __name__ == '__main__':
    response = get(cmd='GetBackLinkData',
                   params={'item': 'http://boingboing.net',
                           'Count': '5',
                           'datasource': 'fresh'})
    print(response)
<commit_msg>Extract params from GetIndexItemInfo.<commit_after>import argparse
import json
import os

import requests

BASE_URL = "https://api.majestic.com/api/json"
BASE_PARAMS = {'app_api_key': os.environ.get('THEAPIKEY')}

def get(cmd, params):
    querydict = {'cmd': cmd}
    querydict.update(BASE_PARAMS)
    querydict.update(params)
    response = requests.get(BASE_URL, params=querydict)
    return json.loads(response.text)

def getIndexItemInfo(site):
    cmd = 'GetIndexItemInfo'
    params = {'items': '2',
              'item0': site,
              'item1': 'chrishannam.co.uk',
              'datasource': 'fresh'}
    responsedata = get(cmd, params)
    if responsedata['Code'] == 'OK':
        data = responsedata['DataTables']['Results']['Data'][0]
        for data in responsedata['DataTables']['Results']['Data']:
            yield {
                'speed': data['OutDomainsExternal'] + 1,
                'power': data['OutLinksExternal'] + 1,
                'agility': data['OutLinksInternal'] + 1,
                'strength': data['RefDomainsEDU'] + 1,
                'smell': data['CitationFlow'] + 1,
            }
    else:
        yield {}

def run():
    parser = argparse.ArgumentParser(description="a test thing")
    parser.add_argument('url')

    args = parser.parse_args()
    results = getIndexItemInfo(args.url)
    for result in results:
        print(result)

if __name__ == '__main__':
    run()
<|endoftext|>


<commit_before>"""Module that contains NaiveRoleAssignment. """

from math import isfinite
import sys
from typing import List, Optional, Tuple

import numpy as np
import scipy.optimize
import stp.rc
import stp.role as role
import stp.role.assignment as assignment
from stp.role import RoleResult

SortedRequests = List[assignment.FlatRoleRequests]

# Define some big constant for "hard constraints".
INVALID_COST = sys.maxsize

class NaiveRoleAssignment(assignment.IRoleAssignment):
    """Naive implementation of role assignment that simply performs the Hungarian
    Algorithm (from scipy.optimize) on HIGH, then MEDIUM, then LOW priority in that
    order."""

    def __init__(self):
        self.prev_assignments = None

    @staticmethod
    def get_sorted_requests(requests: assignment.FlatRoleRequests) -> SortedRequests:
        """Returns a list of FlatRoleRequests sorted in ascending priority order.
        :param requests: Flat list of requests.
        :return: List of FlatRoleRequests in sorted ascending priority order, ie.
        [LOW, MEDIUM, HIGH].
        """
        role_id: assignment.RoleId
        request: role.RoleRequest

        sorted_requests: SortedRequests = [
            {} for _ in range(role.Priority.NUM_PRIORITIES)
        ]

        for role_id, request in requests.items():
            sorted_requests[int(request.priority)][role_id] = request

        return sorted_requests

    @staticmethod
    def compute_costs_matrix(
        free_robots: np.ndarray,
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        prev_results: assignment.FlatRoleResults,
    ) -> np.ndarray:
        """Computes the n_free_robots x n_requests cost matrix corresponding to the
        passed in free robots and role requests.
        :param flat_requests: Role requests to compute cost matrix for.
        :param free_robots: Free robots to compute cost matrix for.
        :param world_state: Current world state.
        :param prev_results: The previous results.
        :return: The m x n cost matrix corresponding to the passed in free robots and
        role requests.
        """

        # Cost matrix size (n_free_robots + 1, n_requests), extra row for unassigned cost
        robot_costs: np.ndarray = np.zeros((len(free_robots) + 1, len(flat_requests)))

        # Row index for unassigned cost of role request
        unassigned_idx: float = len(free_robots)

        # Iterate over each role request.
        request: role.RoleRequest
        for request_idx, (role_id, request) in enumerate(flat_requests.items()):
            # For each role request, iterate over robots.
            for robot_idx, robot in enumerate(free_robots):

                # If the robot is not visible then do not consider its cost
                # Cannot be infinte, will break linear_sum_assignment
                if not robot.visible:
                    robot_costs[robot_idx, request_idx] = 1e9
                    robot_costs[unassigned_idx, request_idx] = 1e9
                    continue

                # Get the previous result for this role_id, if available.
                if prev_results is not None:
                    prev_result: Optional[RoleResult] = prev_results.get(
                        role_id, None)
                else:
                    prev_result = None
                # If the constraints are not satisfied, set the cost to INVALID_COST
                # and continue.
                if not request.constraint_fn(robot, prev_result, world_state):
                    robot_costs[robot_idx, request_idx] = INVALID_COST
                    continue

                # Otherwise, record the cost.
                cost: float = request.cost_fn(robot, prev_result, world_state)

                # Throw an exception if the returned cost is not finite,
                # so unassigned roles can be given infinite cost to be never be assigned
                if not isfinite(cost):
                    raise ValueError(
                        "Got a non-finite cost ({}) for request {} and robot {}".format(
                            cost, request, robot
                        )
                    )

                # Add to robot_costs.
                robot_costs[robot_idx, request_idx] = cost

            # Get cost of not assigning a robot to the request
            unassigned_cost: float = request.cost_fn.unassigned_cost_fn(prev_results, world_state)

            # Throw an exception if the returned cost is not finite.
            if not isfinite(unassigned_cost):
                raise ValueError(
                    "Got a non-finite cost ({}) for request {} and \
                    unassinged robot".format(cost, request, robot))

            # Add unassigned cost to last row of robot_costs
            robot_costs[unassigned_idx, request_idx] = unassigned_cost

        return robot_costs

    @staticmethod
    def assign_prioritized_roles(
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        free_robots: np.ndarray,
        prev_results: assignment.FlatRoleResults,
    ) -> Tuple[assignment.FlatRoleResults, np.ndarray]:
        """Assigns roles using the Hungarian algorithm.
        :param flat_requests: The role requests.
        :param world_state: The current state of the game.
        :param free_robots: The array of free robots that haven't been assigned yet.
        This list will be mutated. Array of stp.rc.Robot.
        :param prev_results: The previous results.
        :return: The results of the role assignment and the new free_robots after
        assignment.
        """

        # Create the flat results.
        flat_results: assignment.FlatRoleResults = {
            role_id: role.RoleResult.from_request(request)
            for role_id, request in flat_requests.items()
        }

        # Create a list of keys, so that we have a mapping from indices -> dict entries.
        keys_list: List[assignment.RoleId] = list(flat_requests.keys())

        # Compute the n x m cost matrix.
        robot_costs: np.ndarray = NaiveRoleAssignment.compute_costs_matrix(
            free_robots, flat_requests, world_state, prev_results
        )

        # Get row index of unassigned cost
        unassigned_idx: float = len(free_robots)

        # Get the optimal assignment using the Hungarian algorithm.
        robot_ind: np.ndarray
        request_ind: np.ndarray
        robot_ind, request_ind = scipy.optimize.linear_sum_assignment(robot_costs)

        # Set the assigned robot for each request.
        for assignment_idx in range(request_ind.shape[0]):
            request_idx: int = request_ind[assignment_idx]
            robot_idx: int = robot_ind[assignment_idx]
            cost: float = robot_costs[robot_idx, request_idx]
            # If cost is equal to INVALID_COST, then it means that role assignment
            # failed. Throw an exception.
            if cost == INVALID_COST:
                raise RuntimeError("Hard constraints not satisfied!")

            # If the row index is not the unassigned row, fill the request
            if robot_idx < unassigned_idx:
                if robot_costs[robot_idx , request_idx] < robot_costs[unassigned_idx, request_idx]:
                    robot: stp.rc.Robot = free_robots[robot_idx]
                    role_id = keys_list[request_idx]
                    flat_results[role_id].assign(robot, cost)

        # If unassigned index chosen, remove it from the list of row indices to prevent removing unassigned index from free_robots
        if unassigned_idx in robot_ind:
            robot_ind = robot_ind[robot_ind != unassigned_idx]

        # Remove the robots from free_robots.
        free_robots = np.delete(free_robots, robot_ind)

        return flat_results, free_robots

    def assign_roles(
        self,
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        prev_results: assignment.FlatRoleResults,
    ) -> assignment.FlatRoleResults:
        """Assigns roles.
        :param flat_requests: The role requests.
        :param world_state: The current state of the game.
        :param prev_results: The previous results.
        :return: The results of the role assignment.
        """
        # Collect high, medium and low priority requests, sorted in ascending priority
        # order.
        sorted_requests: SortedRequests = NaiveRoleAssignment.get_sorted_requests(
            flat_requests
        )

        # Assign roles for each priority from HIGH to LOW.
        requests_dict: assignment.FlatRoleRequests

        # Make a copy of rc.our_robots so that we can mutate it without mutating
        # rc.our_robots.
        free_robots = np.array(world_state.our_robots)

        flat_results: assignment.FlatRoleResults = {}

        # Iterate over requests from HIGH to LOW.
        if world_state is not None:
            for requests_dict in reversed(sorted_requests):
                """ Actually perform the assignment
                using the Hungarian algorithm."""

                (
                    prioritized_results,
                    free_robots,
                ) = NaiveRoleAssignment.assign_prioritized_roles(
                    requests_dict, world_state, free_robots,
                    self.prev_assignments)

                # Add the prioritized_results to flat_results.
                flat_results.update(prioritized_results)
        self.prev_assignments = flat_results

        return flat_results
<commit_msg>fix style again<commit_after>"""Module that contains NaiveRoleAssignment. """

from math import isfinite
import sys
from typing import List, Optional, Tuple

import numpy as np
import scipy.optimize
import stp.rc
import stp.role as role
import stp.role.assignment as assignment
from stp.role import RoleResult

SortedRequests = List[assignment.FlatRoleRequests]

# Define some big constant for "hard constraints".
INVALID_COST = sys.maxsize

class NaiveRoleAssignment(assignment.IRoleAssignment):
    """Naive implementation of role assignment that simply performs the Hungarian
    Algorithm (from scipy.optimize) on HIGH, then MEDIUM, then LOW priority in that
    order."""

    def __init__(self):
        self.prev_assignments = None

    @staticmethod
    def get_sorted_requests(requests: assignment.FlatRoleRequests) -> SortedRequests:
        """Returns a list of FlatRoleRequests sorted in ascending priority order.
        :param requests: Flat list of requests.
        :return: List of FlatRoleRequests in sorted ascending priority order, ie.
        [LOW, MEDIUM, HIGH].
        """
        role_id: assignment.RoleId
        request: role.RoleRequest

        sorted_requests: SortedRequests = [
            {} for _ in range(role.Priority.NUM_PRIORITIES)
        ]

        for role_id, request in requests.items():
            sorted_requests[int(request.priority)][role_id] = request

        return sorted_requests

    @staticmethod
    def compute_costs_matrix(
        free_robots: np.ndarray,
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        prev_results: assignment.FlatRoleResults,
    ) -> np.ndarray:
        """Computes the n_free_robots x n_requests cost matrix corresponding to the
        passed in free robots and role requests.
        :param flat_requests: Role requests to compute cost matrix for.
        :param free_robots: Free robots to compute cost matrix for.
        :param world_state: Current world state.
        :param prev_results: The previous results.
        :return: The m x n cost matrix corresponding to the passed in free robots and
        role requests.
        """

        # Cost matrix size (n_free_robots + 1, n_requests), extra row for unassigned cost
        robot_costs: np.ndarray = np.zeros((len(free_robots) + 1, len(flat_requests)))

        # Row index for unassigned cost of role request
        unassigned_idx: float = len(free_robots)

        # Iterate over each role request.
        request: role.RoleRequest
        for request_idx, (role_id, request) in enumerate(flat_requests.items()):
            # For each role request, iterate over robots.
            for robot_idx, robot in enumerate(free_robots):

                # If the robot is not visible then do not consider its cost
                # Cannot be infinte, will break linear_sum_assignment
                if not robot.visible:
                    robot_costs[robot_idx, request_idx] = 1e9
                    robot_costs[unassigned_idx, request_idx] = 1e9
                    continue

                # Get the previous result for this role_id, if available.
                if prev_results is not None:
                    prev_result: Optional[RoleResult] = prev_results.get(
                        role_id, None)
                else:
                    prev_result = None
                # If the constraints are not satisfied, set the cost to INVALID_COST
                # and continue.
                if not request.constraint_fn(robot, prev_result, world_state):
                    robot_costs[robot_idx, request_idx] = INVALID_COST
                    continue

                # Otherwise, record the cost.
                cost: float = request.cost_fn(robot, prev_result, world_state)

                # Throw an exception if the returned cost is not finite,
                # so unassigned roles can be given infinite cost to be never be assigned
                if not isfinite(cost):
                    raise ValueError(
                        "Got a non-finite cost ({}) for request {} and robot {}".format(
                            cost, request, robot
                        )
                    )

                # Add to robot_costs.
                robot_costs[robot_idx, request_idx] = cost

            # Get cost of not assigning a robot to the request
            unassigned_cost: float = request.cost_fn.unassigned_cost_fn(prev_results, world_state)

            # Throw an exception if the returned cost is not finite.
            if not isfinite(unassigned_cost):
                raise ValueError(
                    "Got a non-finite cost ({}) for request {} and \
                    unassinged robot {}".format(cost, request, robot))

            # Add unassigned cost to last row of robot_costs
            robot_costs[unassigned_idx, request_idx] = unassigned_cost

        return robot_costs

    @staticmethod
    def assign_prioritized_roles(
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        free_robots: np.ndarray,
        prev_results: assignment.FlatRoleResults,
    ) -> Tuple[assignment.FlatRoleResults, np.ndarray]:
        """Assigns roles using the Hungarian algorithm.
        :param flat_requests: The role requests.
        :param world_state: The current state of the game.
        :param free_robots: The array of free robots that haven't been assigned yet.
        This list will be mutated. Array of stp.rc.Robot.
        :param prev_results: The previous results.
        :return: The results of the role assignment and the new free_robots after
        assignment.
        """

        # Create the flat results.
        flat_results: assignment.FlatRoleResults = {
            role_id: role.RoleResult.from_request(request)
            for role_id, request in flat_requests.items()
        }

        # Create a list of keys, so that we have a mapping from indices -> dict entries.
        keys_list: List[assignment.RoleId] = list(flat_requests.keys())

        # Compute the n x m cost matrix.
        robot_costs: np.ndarray = NaiveRoleAssignment.compute_costs_matrix(
            free_robots, flat_requests, world_state, prev_results
        )

        # Get row index of unassigned cost
        unassigned_idx: float = len(free_robots)

        # Get the optimal assignment using the Hungarian algorithm.
        robot_ind: np.ndarray
        request_ind: np.ndarray
        robot_ind, request_ind = scipy.optimize.linear_sum_assignment(robot_costs)

        # Set the assigned robot for each request.
        for assignment_idx in range(request_ind.shape[0]):
            request_idx: int = request_ind[assignment_idx]
            robot_idx: int = robot_ind[assignment_idx]
            cost: float = robot_costs[robot_idx, request_idx]
            # If cost is equal to INVALID_COST, then it means that role assignment
            # failed. Throw an exception.
            if cost == INVALID_COST:
                raise RuntimeError("Hard constraints not satisfied!")

            # If the row index is not the unassigned row, fill the request
            if robot_idx < unassigned_idx:
                if robot_costs[robot_idx , request_idx] < robot_costs[unassigned_idx, request_idx]:
                    robot: stp.rc.Robot = free_robots[robot_idx]
                    role_id = keys_list[request_idx]
                    flat_results[role_id].assign(robot, cost)

        # If unassigned index chosen, remove it from the list of row indices to prevent removing unassigned index from free_robots
        if unassigned_idx in robot_ind:
            robot_ind = robot_ind[robot_ind != unassigned_idx]

        # Remove the robots from free_robots.
        free_robots = np.delete(free_robots, robot_ind)

        return flat_results, free_robots

    def assign_roles(
        self,
        flat_requests: assignment.FlatRoleRequests,
        world_state: stp.rc.WorldState,
        prev_results: assignment.FlatRoleResults,
    ) -> assignment.FlatRoleResults:
        """Assigns roles.
        :param flat_requests: The role requests.
        :param world_state: The current state of the game.
        :param prev_results: The previous results.
        :return: The results of the role assignment.
        """
        # Collect high, medium and low priority requests, sorted in ascending priority
        # order.
        sorted_requests: SortedRequests = NaiveRoleAssignment.get_sorted_requests(
            flat_requests
        )

        # Assign roles for each priority from HIGH to LOW.
        requests_dict: assignment.FlatRoleRequests

        # Make a copy of rc.our_robots so that we can mutate it without mutating
        # rc.our_robots.
        free_robots = np.array(world_state.our_robots)

        flat_results: assignment.FlatRoleResults = {}

        # Iterate over requests from HIGH to LOW.
        if world_state is not None:
            for requests_dict in reversed(sorted_requests):
                """ Actually perform the assignment
                using the Hungarian algorithm."""

                (
                    prioritized_results,
                    free_robots,
                ) = NaiveRoleAssignment.assign_prioritized_roles(
                    requests_dict, world_state, free_robots,
                    self.prev_assignments)

                # Add the prioritized_results to flat_results.
                flat_results.update(prioritized_results)
        self.prev_assignments = flat_results

        return flat_results
<|endoftext|>