<commit_before>"""
###############################################################################
Controller:  Overall controller class
###############################################################################
"""
import pickle as _pickle
import copy as _copy
import time
import random
import string
import OpenPNM
from OpenPNM.Base import logging
logger = logging.getLogger()

class Controller(dict):
    # The following __instance__ class variable and subclassed __new__ method
    # makes the Controller class a 'Singleton'.  This way, the _ctrl attribute
    # of every OpenPNM object is the same, AND if you create a ctrl on the
    # command line (ctrl = OpenPNM.Base.Controller()) it will be the same ctrl!
    __instance__ = None

    def __new__(cls, *args, **kwargs):
        if Controller.__instance__ is None:
            Controller.__instance__ = dict.__new__(cls)
        return Controller.__instance__

    def __init__(self):
        self.comments = 'Using OpenPNM ' + OpenPNM.__version__

    def __str__(self):
        lines = []
        horizontal_rule = 60 * '-'
        for net in self.networks():
            lines.append(horizontal_rule)
            lines.append('{0:<15} {1:<20} ({2})'.format('Object:',
                                                        'Name',
                                                        'Class'))
            lines.append(horizontal_rule)
            lines.append('{0:<15} {1:<20} ({2})'.format('Network:',
                                                        net.name,
                                                        net.__class__.__name__))
            for geom in net._geometries:
                str = '++ {0:<12} {1:<20} ({2})'
                lines.append(str.format('Geometry: ',
                                        geom.name,
                                        geom.__class__.__name__))
            for phase in net._phases:
                if len(phase._phases) == 0:
                    str = '+ {0:<13} {1:<20} ({2})'
                    lines.append(str.format('Pure Phase: ',
                                            phase.name,
                                            phase.__class__.__name__))
                if len(phase._phases) > 1:
                    str = '+ {0:<13} {1:<20} ({2})'
                    lines.append(str.format('Mixture Phase: ',
                                            phase.name,
                                            phase.__class__.__name__))
                    comps = phase.phases()
                    for compname in comps:
                        str = '++ {0:<12} {1:<20} ({2})}'
                        lines.append(str.format('Component Phase: ',
                                                phase.name,
                                                phase.__class__.__name__))
                for phys in phase._physics:
                    str = '++ {0:<12} {1:<20} ({2})'
                    lines.append(str.format('Physics: ',
                                            phys.name,
                                            phys.__class__.__name__))
        return '\n'.join(lines)

    def _setloglevel(self, level):
        logger.setLevel(level)

    def _getloglevel(self):
        return 'Log level is currently set to: ' + str(logger.level)

    loglevel = property(fget=_getloglevel, fset=_setloglevel)

    def networks(self):
        r"""
        Returns a list of all Network objects
        """
        return self._get_objects(obj_type='GenericNetwork')

    def geometries(self):
        r"""
        Returns a list of all Geometry objects
        """
        return self._get_objects(obj_type='GenericGeometry')

    def phases(self):
        r"""
        Returns a list of all Phase objects
        """
        return self._get_objects(obj_type='GenericPhase')

    def physics(self):
        r"""
        Returns a list of all Physics objects
        """
        return self._get_objects(obj_type='GenericPhysics')

    def algorithms(self):
        r"""
        Returns a list of all Algorithm objects
        """
        return self._get_objects(obj_type='GenericAlgorithm')

    def _get_objects(self, obj_type):
        temp = []
        for obj in self.keys():
            mro = [item.__name__ for item in self[obj].__class__.__mro__]
            if obj_type in mro:
                temp.append(self[obj])
        return temp

    def clear(self):
        r"""
        This is an overloaded version of the standard dict's ``clear`` method.
        This completely clears the Controller object's dict as expected, but
        also removes links to the Controller object in all objects.

        """
        for item in self.keys():
            self[item]._ctrl = {}
        self.__dict__ = {}
        super().clear()

    def update(self, arg):
        r"""
        This is a subclassed version of the standard dict's ``update`` method.
        It can accept a dictionary of OpenPNM Core objects in which case it
        adds the objects to the Controller.  It can also accept an OpenPNM
        Network object in which case it extracts all associated objects and
        adds them to the Controller.  In both cases it adds the Controller to
        all object's ``controller`` attribute.

        Notes
        -----
        The Network (and other Core objects) do not store Algorithms, so this
        update will not add any Algorithm objects to the Controller.  This may
        change.
        """
        if arg.__class__ == dict:
            for item in arg.keys():
                self[item] = arg[item]
                arg[item]._ctrl = self
        else:
            mro = [item.__name__ for item in arg.__class__.__mro__]
            if 'GenericNetwork' in mro:
                net = arg
                self[net.name] = net
                net._ctrl = self
                for item in net._geometries + net._physics + net._phases:
                    self[item.name] = item
                    item._ctrl = self

    def purge_object(self, obj, mode='single'):
        r"""
        Remove an object, including all traces of it in its associated objects

        Parameters
        ----------
        obj : OpenPNM Object
            The object to be removed.  This method removes all traces of the
            object from everywhere, including all the object tracking lists and
            label dictionaries of every object.
        mode : string
            Dicates the type of purge to be performed.  Options are:

            - 'single': Only purges the specified object
            - 'complete': Purges the specified object AND all of its associated
                          objects

        Notes
        -----
        To only remove an object from the Contoller object use the dictionary's
        native ``pop`` method.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> geom = OpenPNM.Geometry.GenericGeometry(network=pn,
        ...                                         pores=pn.Ps,
        ...                                         throats=pn.Ts)

        # Label entries are added to the Network where geom is defined
        >>> 'pore.'+geom.name in pn.keys()
        True
        >>> ctrl.purge_object(geom)

        # geom is removed from Controller object
        >>> geom.name in ctrl.keys()
        False

        # geom's labels are removed from the Network too
        >>> 'pore.' + geom.name in pn.keys()
        False
        """
        if mode == 'complete':
            if obj._net is None:
                net = obj
            else:
                net = obj._net
            for item in net.geometries() + net.phases() + net.physics():
                blank = self.pop(item, None)
            del self[net.name]
        elif mode == 'single':
            name = obj.name
            for item in list(self.keys()):
                # Remove label arrays from all other objects
                self[item].pop('pore.' + name, None)
                self[item].pop('throat.' + name, None)
                # Remove associations on other objects
                self[item]._geometries[:] = \
                    [x for x in self[item]._geometries if x is not obj]
                self[item]._phases[:] = \
                    [x for x in self[item]._phases if x is not obj]
                self[item]._physics[:] = \
                    [x for x in self[item]._physics if x is not obj]
            # Set object's controller attribute to an empty dict
            self[name]._ctrl = {}
            # Remove object from Controller dict
            del self[name]

    def ghost_object(self, obj):
        r"""
        Create a ghost OpenPNM Object containing all the data, methods and
        associations of the original object, but without registering the ghost
        anywhere.   This ghost is intended as a disposable object, for
        instance, to receive data without overwriting existing data.

        Parameters
        ----------
        obj : OpenPNM Object
            The object to be cloned can be any OpenPNM Object

        Returns
        -------
        A clone of the specified object is returned, but it retains all its links
        to the objects associated with the original object.  The cloned object is
        not associated with the Network.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> pn2 = ctrl.ghost_object(pn)
        >>> pn is pn2  # A copy of pn is created
        False
        >>> pn2.keys() == pn.keys()  # They have otherwise identical data
        True
        >>> pn2.controller is ctrl # pn2 is not associated with existing Controller
        False

        It can also be used to create ghosts of other object types:

        >>> geom = OpenPNM.Geometry.TestGeometry(network=pn,
        ...                                      pores=pn.Ps,
        ...                                      throats=pn.Ts)
        >>> geo2 = ctrl.ghost_object(geom)
        >>> geom is geo2
        False

        # Ghost has same name as ancestor
        >>> geom.name == geo2.name
        True

        # But they are not the same object
        >>> geo2 is ctrl[geo2.name]
        False

        # The ghost is not registered with the Controller
        >>> geo2.controller
        {}

        # The following comparisons look at some 'behind the scenes' information
        # The ghost and ancestor are assoicated with the same Network
        >>> geo2._net == geom._net
        True

        # But the Network remains aware of the ancestor only
        >>> geo2 in pn._geometries
        False

        """
        obj_new = _copy.copy(obj)
        obj_new.__dict__ = obj.__dict__
        obj_new._ctrl = {}
        del self[obj.name]
        self[obj.name] = obj
        return obj_new

    def save_simulation(self, network, filename=''):
        r"""
        Save a single Network simulation to a 'net' file, including all of its
        associated objects, but not Algorithms

        Parameters
        ----------
        network : OpenPNM Network object
            The Network to save
        filename : string, optional
            If no filename is given the name of the Network is used
        """
        if filename == '':
            filename = network.name
        else:
            filename = filename.rstrip('.net')

        network._ctrl = {}
        # Save nested dictionary pickle
        _pickle.dump(network, open(filename + '.net', 'wb'))
        network._ctrl = self

    def load_simulation(self, filename):
        r"""
        Loads a Network simulation fromt the specified 'net' file and adds it
        to the Controller

        Parameters
        ----------
        filename : string
            The name of the file containing the Network simulation to load
        """
        filename = filename.rstrip('.net')
        net = _pickle.load(open(filename + '.net', 'rb'))
        net.controller = self

    def save(self, filename=''):
        r"""
        Save the entire state of the Controller to a 'pnm' file.

        Parameters
        ----------
        filename : string, optional
            The file name to save as. If no filename is provided the current
            date and time is used.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> ctrl.save('test.pnm')
        >>> pn.name in ctrl.keys()
        True
        >>> ctrl.clear()
        >>> ctrl.keys()
        dict_keys([])
        >>> ctrl.load('test.pnm')
        >>> pn.name in ctrl.keys()
        True
        """
        if filename == '':
            from datetime import datetime
            i = datetime.now()
            filename = i.strftime('%Y-%m-%d_%H-%M-%S')
        else:
            filename = filename.rstrip('.pnm')

        # Save nested dictionary pickle
        _pickle.dump(self, open(filename + '.pnm', 'wb'))

    def load(self, filename):
        r"""
        Load an entire Controller from a 'pnm' file.

        Parameters
        ----------
        filename : string
            The file name of the Controller to load.

        Notes
        -----
        This calls the ``clear`` method of the Controller object, so it will
        over write the calling objects information AND remove any references
        to the calling object from existing objects.
        """
        filename = filename.strip('.pnm')
        if self != {}:
            print('Warning: Loading data onto non-empty controller object, \
                   existing data will be lost')
            self.clear()
        self = _pickle.load(open(filename + '.pnm', 'rb'))

    def export(self, network=None, filename='', fileformat='VTK'):
        r"""
        Export data to the specified file format.

        Parameters
        ----------
        network : OpenPNM Network Object
            This Network and all of its phases will be written to the specified
            file.  If no Netowrk is given it will check to ensure that only one
            Network exists on the Controller and use that.  If there is more
            than one Network an error is thrown.
        filename : string, optional
            The file name to save as.  If no name is given then the name of
            suppiled object is used.  If no object is given, the name of the
            Network is used.
        fileformat : string
            The type of file to create.  Options are:

            1. VTK: Suitable for visualizing in VTK capable software such as Paraview
            2. MAT: Suitable for loading data into Matlab for post-processing

        """
        if network is None:
            if len(self.networks()) == 1:
                network = self.networks()[0]
            else:
                raise Exception('Multiple Networks found, please specify \
                                which to export')
        import OpenPNM.Utilities.IO as io
        if fileformat == 'VTK':
            phases = network._phases
            io.VTK.save(filename=filename, network=network, phases=phases)
            return
        if fileformat == 'MAT':
            phases = network._phases
            io.MAT.save(filename=filename, network=network, phases=phases)
            return

    def _script(self, filename, mode='read'):
        r"""
        Save or reload the script files used for the modeling

        Parameters
        ----------
        filename : string
            The name of the file to read or write
        mode : string
            Whether to 'archive' the given script file on the object or to
            'retrieve' it from the object and create a new file with it.  The
            default is 'archive'.
        """
        filename = filename.split('.')[0]+'.py'
        if mode == 'archive':
            with open(filename, 'rb') as read_file:
                contents = read_file.read()
            self._script = contents
        if mode == 'retrieve':
            with open(filename, 'wb') as write_file:
                write_file.write(self._script)

    def _set_comments(self, string):
        if hasattr(self, '_comments') is False:
            self._comments = {}
        self._comments[time.strftime('%c')] = string

    def _get_comments(self):
        if hasattr(self, '_comments') is False:
            print('No comments found')
        else:
            for key in self._comments.keys():
                print(key, ': ', self._comments[key])

    comments = property(fget=_get_comments, fset=_set_comments)

    def clone_simulation(self, network, name=None):
        r"""
        Accepts a Network object and creates a complete clone including all
        associated objects.  All objects in the cloned simulation are
        registered with the Controller object and are fully functional.

        Returns
        -------
        A handle to the new Network object, which will include handles to
        clones of all associated objects.

        See Also
        --------
        ghost_object

        Notes
        -----
        One useful application of this method is to create a cloned simulation
        that can be trimmed to a smaller size.  This small simulation will
        result in much faster Algorithms calculations.

        Examples
        --------
        None yet
        """
        if network._parent is not None:
            logger.error('Cannot clone a network that is already a clone')
            return
        bak = {}
        bak.update(self)
        self.clear()
        net = _copy.deepcopy(network)
        self.update(net)
        if name is None:
            for _ in range(5):
                name = ''.join(random.choice(string.ascii_uppercase +
                                             string.ascii_lowercase +
                                             string.digits))
        for item in list(self.keys()):
            temp = self.pop(item)
            temp._parent = network
            new_name = temp.name + '_' + name
            temp.name = new_name
            self[temp.name] = temp
        # Add parent Network numbering to clone
        net['pore.' + network.name] = network.Ps
        net['throat.' + network.name] = network.Ts
        self.update(bak)
        return net

if __name__ == '__main__':
    ctrl = Controller()
<commit_msg>Fixed ghost object bug<commit_after>"""
###############################################################################
Controller:  Overall controller class
###############################################################################
"""
import pickle as _pickle
import copy as _copy
import time
import random
import string
import OpenPNM
from OpenPNM.Base import logging
logger = logging.getLogger()

class Controller(dict):
    # The following __instance__ class variable and subclassed __new__ method
    # makes the Controller class a 'Singleton'.  This way, the _ctrl attribute
    # of every OpenPNM object is the same, AND if you create a ctrl on the
    # command line (ctrl = OpenPNM.Base.Controller()) it will be the same ctrl!
    __instance__ = None

    def __new__(cls, *args, **kwargs):
        if Controller.__instance__ is None:
            Controller.__instance__ = dict.__new__(cls)
        return Controller.__instance__

    def __init__(self):
        self.comments = 'Using OpenPNM ' + OpenPNM.__version__

    def __str__(self):
        lines = []
        horizontal_rule = 60 * '-'
        for net in self.networks():
            lines.append(horizontal_rule)
            lines.append('{0:<15} {1:<20} ({2})'.format('Object:',
                                                        'Name',
                                                        'Class'))
            lines.append(horizontal_rule)
            lines.append('{0:<15} {1:<20} ({2})'.format('Network:',
                                                        net.name,
                                                        net.__class__.__name__))
            for geom in net._geometries:
                str = '++ {0:<12} {1:<20} ({2})'
                lines.append(str.format('Geometry: ',
                                        geom.name,
                                        geom.__class__.__name__))
            for phase in net._phases:
                if len(phase._phases) == 0:
                    str = '+ {0:<13} {1:<20} ({2})'
                    lines.append(str.format('Pure Phase: ',
                                            phase.name,
                                            phase.__class__.__name__))
                if len(phase._phases) > 1:
                    str = '+ {0:<13} {1:<20} ({2})'
                    lines.append(str.format('Mixture Phase: ',
                                            phase.name,
                                            phase.__class__.__name__))
                    comps = phase.phases()
                    for compname in comps:
                        str = '++ {0:<12} {1:<20} ({2})}'
                        lines.append(str.format('Component Phase: ',
                                                phase.name,
                                                phase.__class__.__name__))
                for phys in phase._physics:
                    str = '++ {0:<12} {1:<20} ({2})'
                    lines.append(str.format('Physics: ',
                                            phys.name,
                                            phys.__class__.__name__))
        return '\n'.join(lines)

    def _setloglevel(self, level):
        logger.setLevel(level)

    def _getloglevel(self):
        return 'Log level is currently set to: ' + str(logger.level)

    loglevel = property(fget=_getloglevel, fset=_setloglevel)

    def networks(self):
        r"""
        Returns a list of all Network objects
        """
        return self._get_objects(obj_type='GenericNetwork')

    def geometries(self):
        r"""
        Returns a list of all Geometry objects
        """
        return self._get_objects(obj_type='GenericGeometry')

    def phases(self):
        r"""
        Returns a list of all Phase objects
        """
        return self._get_objects(obj_type='GenericPhase')

    def physics(self):
        r"""
        Returns a list of all Physics objects
        """
        return self._get_objects(obj_type='GenericPhysics')

    def algorithms(self):
        r"""
        Returns a list of all Algorithm objects
        """
        return self._get_objects(obj_type='GenericAlgorithm')

    def _get_objects(self, obj_type):
        temp = []
        for obj in self.keys():
            mro = [item.__name__ for item in self[obj].__class__.__mro__]
            if obj_type in mro:
                temp.append(self[obj])
        return temp

    def clear(self):
        r"""
        This is an overloaded version of the standard dict's ``clear`` method.
        This completely clears the Controller object's dict as expected, but
        also removes links to the Controller object in all objects.

        """
        for item in self.keys():
            self[item]._ctrl = {}
        self.__dict__ = {}
        super().clear()

    def update(self, arg):
        r"""
        This is a subclassed version of the standard dict's ``update`` method.
        It can accept a dictionary of OpenPNM Core objects in which case it
        adds the objects to the Controller.  It can also accept an OpenPNM
        Network object in which case it extracts all associated objects and
        adds them to the Controller.  In both cases it adds the Controller to
        all object's ``controller`` attribute.

        Notes
        -----
        The Network (and other Core objects) do not store Algorithms, so this
        update will not add any Algorithm objects to the Controller.  This may
        change.
        """
        if arg.__class__ == dict:
            for item in arg.keys():
                self[item] = arg[item]
                arg[item]._ctrl = self
        else:
            mro = [item.__name__ for item in arg.__class__.__mro__]
            if 'GenericNetwork' in mro:
                net = arg
                self[net.name] = net
                net._ctrl = self
                for item in net._geometries + net._physics + net._phases:
                    self[item.name] = item
                    item._ctrl = self

    def purge_object(self, obj, mode='single'):
        r"""
        Remove an object, including all traces of it in its associated objects

        Parameters
        ----------
        obj : OpenPNM Object
            The object to be removed.  This method removes all traces of the
            object from everywhere, including all the object tracking lists and
            label dictionaries of every object.
        mode : string
            Dicates the type of purge to be performed.  Options are:

            - 'single': Only purges the specified object
            - 'complete': Purges the specified object AND all of its associated
                          objects

        Notes
        -----
        To only remove an object from the Contoller object use the dictionary's
        native ``pop`` method.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> geom = OpenPNM.Geometry.GenericGeometry(network=pn,
        ...                                         pores=pn.Ps,
        ...                                         throats=pn.Ts)

        # Label entries are added to the Network where geom is defined
        >>> 'pore.'+geom.name in pn.keys()
        True
        >>> ctrl.purge_object(geom)

        # geom is removed from Controller object
        >>> geom.name in ctrl.keys()
        False

        # geom's labels are removed from the Network too
        >>> 'pore.' + geom.name in pn.keys()
        False
        """
        if mode == 'complete':
            if obj._net is None:
                net = obj
            else:
                net = obj._net
            for item in net.geometries() + net.phases() + net.physics():
                blank = self.pop(item, None)
            del self[net.name]
        elif mode == 'single':
            name = obj.name
            for item in list(self.keys()):
                # Remove label arrays from all other objects
                self[item].pop('pore.' + name, None)
                self[item].pop('throat.' + name, None)
                # Remove associations on other objects
                self[item]._geometries[:] = \
                    [x for x in self[item]._geometries if x is not obj]
                self[item]._phases[:] = \
                    [x for x in self[item]._phases if x is not obj]
                self[item]._physics[:] = \
                    [x for x in self[item]._physics if x is not obj]
            # Set object's controller attribute to an empty dict
            self[name]._ctrl = {}
            # Remove object from Controller dict
            del self[name]

    def ghost_object(self, obj):
        r"""
        Create a ghost OpenPNM Object containing all the data, methods and
        associations of the original object, but without registering the ghost
        anywhere.   This ghost is intended as a disposable object, for
        instance, to receive data without overwriting existing data.

        Parameters
        ----------
        obj : OpenPNM Object
            The object to be cloned can be any OpenPNM Object

        Returns
        -------
        A clone of the specified object is returned, but it retains all its links
        to the objects associated with the original object.  The cloned object is
        not associated with the Network.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> pn2 = ctrl.ghost_object(pn)
        >>> pn is pn2  # A copy of pn is created
        False
        >>> pn2.keys() == pn.keys()  # They have otherwise identical data
        True
        >>> pn2.controller is ctrl # pn2 is not associated with existing Controller
        False

        It can also be used to create ghosts of other object types:

        >>> geom = OpenPNM.Geometry.TestGeometry(network=pn,
        ...                                      pores=pn.Ps,
        ...                                      throats=pn.Ts)
        >>> geo2 = ctrl.ghost_object(geom)
        >>> geom is geo2
        False

        # Ghost has same name as ancestor
        >>> geom.name == geo2.name
        True

        # But they are not the same object
        >>> geo2 is ctrl[geo2.name]
        False

        # The ghost is not registered with the Controller
        >>> geo2.controller
        {}

        # The following comparisons look at some 'behind the scenes' information
        # The ghost and ancestor are assoicated with the same Network
        >>> geo2._net == geom._net
        True

        # But the Network remains aware of the ancestor only
        >>> geo2 in pn._geometries
        False

        """
        obj_new = _copy.copy(obj)
        obj_new.__dict__ = _copy.copy(obj.__dict__)
        obj_new._ctrl = {}
        del self[obj.name]
        self[obj.name] = obj
        return obj_new

    def save_simulation(self, network, filename=''):
        r"""
        Save a single Network simulation to a 'net' file, including all of its
        associated objects, but not Algorithms

        Parameters
        ----------
        network : OpenPNM Network object
            The Network to save
        filename : string, optional
            If no filename is given the name of the Network is used
        """
        if filename == '':
            filename = network.name
        else:
            filename = filename.rstrip('.net')

        network._ctrl = {}
        # Save nested dictionary pickle
        _pickle.dump(network, open(filename + '.net', 'wb'))
        network._ctrl = self

    def load_simulation(self, filename):
        r"""
        Loads a Network simulation fromt the specified 'net' file and adds it
        to the Controller

        Parameters
        ----------
        filename : string
            The name of the file containing the Network simulation to load
        """
        filename = filename.rstrip('.net')
        net = _pickle.load(open(filename + '.net', 'rb'))
        net.controller = self

    def save(self, filename=''):
        r"""
        Save the entire state of the Controller to a 'pnm' file.

        Parameters
        ----------
        filename : string, optional
            The file name to save as. If no filename is provided the current
            date and time is used.

        Examples
        --------
        >>> import OpenPNM
        >>> ctrl = OpenPNM.Base.Controller()
        >>> pn = OpenPNM.Network.TestNet()
        >>> ctrl.save('test.pnm')
        >>> pn.name in ctrl.keys()
        True
        >>> ctrl.clear()
        >>> ctrl.keys()
        dict_keys([])
        >>> ctrl.load('test.pnm')
        >>> pn.name in ctrl.keys()
        True
        """
        if filename == '':
            from datetime import datetime
            i = datetime.now()
            filename = i.strftime('%Y-%m-%d_%H-%M-%S')
        else:
            filename = filename.rstrip('.pnm')

        # Save nested dictionary pickle
        _pickle.dump(self, open(filename + '.pnm', 'wb'))

    def load(self, filename):
        r"""
        Load an entire Controller from a 'pnm' file.

        Parameters
        ----------
        filename : string
            The file name of the Controller to load.

        Notes
        -----
        This calls the ``clear`` method of the Controller object, so it will
        over write the calling objects information AND remove any references
        to the calling object from existing objects.
        """
        filename = filename.strip('.pnm')
        if self != {}:
            print('Warning: Loading data onto non-empty controller object, \
                   existing data will be lost')
            self.clear()
        self = _pickle.load(open(filename + '.pnm', 'rb'))

    def export(self, network=None, filename='', fileformat='VTK'):
        r"""
        Export data to the specified file format.

        Parameters
        ----------
        network : OpenPNM Network Object
            This Network and all of its phases will be written to the specified
            file.  If no Netowrk is given it will check to ensure that only one
            Network exists on the Controller and use that.  If there is more
            than one Network an error is thrown.
        filename : string, optional
            The file name to save as.  If no name is given then the name of
            suppiled object is used.  If no object is given, the name of the
            Network is used.
        fileformat : string
            The type of file to create.  Options are:

            1. VTK: Suitable for visualizing in VTK capable software such as Paraview
            2. MAT: Suitable for loading data into Matlab for post-processing

        """
        if network is None:
            if len(self.networks()) == 1:
                network = self.networks()[0]
            else:
                raise Exception('Multiple Networks found, please specify \
                                which to export')
        import OpenPNM.Utilities.IO as io
        if fileformat == 'VTK':
            phases = network._phases
            io.VTK.save(filename=filename, network=network, phases=phases)
            return
        if fileformat == 'MAT':
            phases = network._phases
            io.MAT.save(filename=filename, network=network, phases=phases)
            return

    def _script(self, filename, mode='read'):
        r"""
        Save or reload the script files used for the modeling

        Parameters
        ----------
        filename : string
            The name of the file to read or write
        mode : string
            Whether to 'archive' the given script file on the object or to
            'retrieve' it from the object and create a new file with it.  The
            default is 'archive'.
        """
        filename = filename.split('.')[0]+'.py'
        if mode == 'archive':
            with open(filename, 'rb') as read_file:
                contents = read_file.read()
            self._script = contents
        if mode == 'retrieve':
            with open(filename, 'wb') as write_file:
                write_file.write(self._script)

    def _set_comments(self, string):
        if hasattr(self, '_comments') is False:
            self._comments = {}
        self._comments[time.strftime('%c')] = string

    def _get_comments(self):
        if hasattr(self, '_comments') is False:
            print('No comments found')
        else:
            for key in self._comments.keys():
                print(key, ': ', self._comments[key])

    comments = property(fget=_get_comments, fset=_set_comments)

    def clone_simulation(self, network, name=None):
        r"""
        Accepts a Network object and creates a complete clone including all
        associated objects.  All objects in the cloned simulation are
        registered with the Controller object and are fully functional.

        Returns
        -------
        A handle to the new Network object, which will include handles to
        clones of all associated objects.

        See Also
        --------
        ghost_object

        Notes
        -----
        One useful application of this method is to create a cloned simulation
        that can be trimmed to a smaller size.  This small simulation will
        result in much faster Algorithms calculations.

        Examples
        --------
        None yet
        """
        if network._parent is not None:
            logger.error('Cannot clone a network that is already a clone')
            return
        bak = {}
        bak.update(self)
        self.clear()
        net = _copy.deepcopy(network)
        self.update(net)
        if name is None:
            for _ in range(5):
                name = ''.join(random.choice(string.ascii_uppercase +
                                             string.ascii_lowercase +
                                             string.digits))
        for item in list(self.keys()):
            temp = self.pop(item)
            temp._parent = network
            new_name = temp.name + '_' + name
            temp.name = new_name
            self[temp.name] = temp
        # Add parent Network numbering to clone
        net['pore.' + network.name] = network.Ps
        net['throat.' + network.name] = network.Ts
        self.update(bak)
        return net

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


<commit_before>
from PLC.Faults import *
from PLC.Method import Method
from PLC.Parameter import Parameter, Mixed
from PLC.Nodes import Node, Nodes
from PLC.NodeGroups import NodeGroup, NodeGroups
#from PLC.Sites import Site, Sites
from PLC.Auth import PasswordAuth

class AdmAddNodeGroup(Method):
    """
    Adds a new node group. Any values specified in optional_vals are used,
    otherwise defaults are used.

    Returns the new node_id (> 0) if successful, faults otherwise.
    """

    roles = ['admin']

    can_update = lambda (field, value): field in \
                 ['name', 'description', 'is_custom']
    update_fields = dict(filter(can_update, NodeGroup.fields.items()))

    accepts = [
        PasswordAuth(),
        NodeGroup.fields['name'],
        NodeGroup.fields['description'],
        update_fields
        ]

    returns = Parameter(int, '1 if successful')

    def call(self, auth, name, description, optional_vals = {}):
        if filter(lambda field: field not in self.update_fields, optional_vals):
            raise PLCInvalidArgument, "Invalid fields specified"

        # Authenticated function
        assert self.caller is not None

        # make sure we are 'admin'
        if 'admin' not in self.caller['roles']:
        	raise PLCPermissionDenied, "Not allowed to add node groups"

	#creat node group
        node_group = NodeGroup(self.api, optional_vals)
        node_group['name'] = name
        node_group['description'] = description
        node_group.flush()

        return node_group['nodegroup_id']
<commit_msg>- fix documentation for return value - remove redundant roles check<commit_after>
from PLC.Faults import *
from PLC.Method import Method
from PLC.Parameter import Parameter, Mixed
from PLC.Nodes import Node, Nodes
from PLC.NodeGroups import NodeGroup, NodeGroups
#from PLC.Sites import Site, Sites
from PLC.Auth import PasswordAuth

class AdmAddNodeGroup(Method):
    """
    Adds a new node group. Any values specified in optional_vals are used,
    otherwise defaults are used.

    Returns the new nodegroup_id (> 0) if successful, faults otherwise.
    """

    roles = ['admin']

    can_update = lambda (field, value): field in \
                 ['name', 'description', 'is_custom']
    update_fields = dict(filter(can_update, NodeGroup.fields.items()))

    accepts = [
        PasswordAuth(),
        NodeGroup.fields['name'],
        NodeGroup.fields['description'],
        update_fields
        ]

    returns = Parameter(int, 'New nodegroup_id (> 0) if successful')

    def call(self, auth, name, description, optional_vals = {}):
        if filter(lambda field: field not in self.update_fields, optional_vals):
            raise PLCInvalidArgument, "Invalid fields specified"

	# Create node group
        node_group = NodeGroup(self.api, optional_vals)
        node_group['name'] = name
        node_group['description'] = description
        node_group.flush()

        return node_group['nodegroup_id']
<|endoftext|>


<commit_before>"""
Находит первое вхождение введённой последовательности в последовательность,
полученную склейкой всех натуральных чисел
"""
from string import *

first100 = ''
for i in range(1, 100005):
    first100 += str(i)
print(first100)

def num_of_crosses(str):
    """
    Возвращает количество 'x' в начале или конце строки
    """
    num = 0
    for char in str:
        if char == 'x':
            num += 1
        else:
            break
    for char in str[::-1]:
        if char == 'x':
            num += 1
        else:
            break
    return num

def num_of_zero(str):
    """
    Возвращает количество нулей в начале строки
    """
    num = 0
    for char in str:
        if char == '0':
            num += 1
        else:
            break
    return num

def fill_crosses(segment, num):
    crosses = num_of_crosses(segment)
    full_num = num.zfill(crosses)
    split_chars = [char for char in segment]
    split_num = [char for char in full_num]
    if segment[0] is 'x':
        for i in range(0, crosses):
            split_chars[i] = split_num[i]
    else:
        split_chars.reverse()
        split_num.reverse()
        for i in range(0, crosses):
            split_chars[i] = split_num[i]
        split_chars.reverse()
    new_segment = ''
    for element in split_chars:
        new_segment += element
    return new_segment

def min_num_to_fill(num_of_crosses, side):
    if num_of_crosses is 0 or side is 'r':
        return 0
    else:
        return int('1'+'0'*(num_of_crosses-1))

def max_num_to_fill(num_of_crosses):
    if num_of_crosses is 0:
        return 0
    return int('9'*num_of_crosses)

def is_empty(segment):
    """
    Проверяет, есть ли в сегменте цифры
    """
    for char in segment:
        if char != 'x':
            return False
    return True

def to_segment(str_seq, segment_len, shift):
    """
    Делит строку на list из значащих сегментов
    """
    ext_str_seq = segment_len*'x'+str_seq+segment_len*'x'
    segments = [ext_str_seq[num+shift:num+segment_len+shift]
                for num in range(0, len(ext_str_seq), segment_len)
                if not is_empty(ext_str_seq[num+shift:num+segment_len+shift])
                ]
    return segments

def get_rank_size(rank):
    if rank is 0:
        return 0
    size = '8'
    for i in range(1, rank):
        size += '9'
    return int(size)

def find_distance(num):
    summ = 0
    num_len = len(str(num))
    for rank in range(1, num_len):
        summ += rank * get_rank_size(rank)
        num -= get_rank_size(rank)
    summ += num*num_len
    summ -= num_len
    return summ

def are_equivalent(num, segment):
    if num == segment:
        return True
    digit_list = [char for char in segment]
    if len(digit_list) is 0:
        return True
    num_list = [char for char in num]
    if len(num_list) < len(digit_list):
        return False
    if segment[-1] is 'x':
        for i in range(0, len(digit_list)-1):
            if num_list[i] != digit_list[i] and digit_list[i] != 'x':
                return False
    else:
        num_list.reverse()
        digit_list.reverse()
        for i in range(0, len(digit_list)):
            if num_list[i] != digit_list[i] and digit_list[i] != 'x':
                return False
    return True

def min_num_in_segments(segments):
    if segments[0][0] is '0':
        return [float('Inf'), 0]
    if len(segments) is 1:
        if segments[0][0] is 'x':
            return [
                fill_crosses(segments[0],
                                str(min_num_to_fill(num_of_crosses(segments[0]), 'l'))),
                num_of_crosses(segments[0])]
        elif segments[0][-1] is 'x':
            return [
                fill_crosses(segments[0],
                                str(min_num_to_fill(num_of_crosses(segments[0]), 'r'))),
                num_of_crosses(segments[0])]
        return [segments[0],
                num_of_crosses(segments[0])]

    full_segments = [segment for segment in segments
                     if not 'x' in segment and len(segment) != 0]
    for segment in full_segments:
        if segment[0] is '0':
            return [float('Inf'), 0]
    if len(full_segments) > 1:
        for i in range(0, len(full_segments)-1):
            if int(full_segments[i])+1 != int(full_segments[i+1]):
                return [float('Inf'), 0]
    if len(full_segments) == len(segments):
        return [segments[0],
                num_of_crosses(segments[0])]

    if segments[0][0] is 'x':
        for i in range(min_num_to_fill(num_of_crosses(segments[0]), 'l'),
                       max_num_to_fill(num_of_crosses(segments[0]))):
            filled = fill_crosses(segments[0], str(i))
            correct = True
            for j in range(1, len(segments)):
                test = (str(int(filled)+j), segments[j])
                if not are_equivalent(str(int(filled)+j), segments[j]):
                    correct = False
            if correct is True:
                return [filled,
                        num_of_crosses(segments[0])]
    if segments[-1][-1] is 'x':
        for i in range(min_num_to_fill(num_of_crosses(segments[-1]), 'r'),
                       max_num_to_fill(num_of_crosses(segments[-1]))):
            filled = fill_crosses(segments[-1], str(i))
            correct = True
            for j in range(1, len(segments)):
                if not are_equivalent(str(int(filled)-j), segments[-j+1]):
                    correct = False
            if correct is True:
                return [filled,
                        num_of_crosses(segments[0])]
    return [float('Inf'), 0]

def split_seq(str_seq):
    nums = dict()
    for segment_len in range(1, len(str_seq)+1+num_of_zero(str_seq)):
        for shift in range(0, segment_len):
            segments = to_segment(str_seq, segment_len, shift)
            min_num = min_num_in_segments(segments)
            if min_num[0] != float('Inf'):
                nums[int(min_num[0])] = min_num[0]
            print(segments)
            print(min_num_in_segments(segments))
    return nums[min(nums.keys())]
#print(split_seq('125'))
# #str_seq = input('Введите искомую последовательность: ')
print(first100.find('151617'))
k = 1
while k < 100005:
    for i in range(k-1, k+2):
        print(i, end='  →  ')
        print(find_distance(i), end='  ~  ')
        print(first100.find(str(i)+str(i+1)+str(i+2)), end='  so  ')
        print('diff is: '+str(find_distance(i) -
                              first100.find(str(i)+str(i+1)+str(i+2))
                              ))
    k *= 10
<commit_msg>v0.3  Заработало, считает корректно, пора прикрутить интерфейс<commit_after>"""
Находит первое вхождение введённой последовательности в последовательность,
полученную склейкой всех натуральных чисел
"""
from string import *

first100 = ''
for i in range(1, 10005):
    first100 += str(i)
print(first100)

def num_of_crosses(str):
    """
    Возвращает количество 'x' в начале или конце строки
    """
    num = 0
    for char in str:
        if char == 'x':
            num += 1
        else:
            break
    for char in str[::-1]:
        if char == 'x':
            num += 1
        else:
            break
    return num

def num_of_zero(str):
    """
    Возвращает количество нулей в начале строки
    """
    num = 0
    for char in str:
        if char == '0':
            num += 1
        else:
            break
    return num

def fill_crosses(segment, num):
    crosses = num_of_crosses(segment)
    full_num = num.zfill(crosses)
    split_chars = [char for char in segment]
    split_num = [char for char in full_num]
    if segment[0] is 'x':
        for i in range(0, crosses):
            split_chars[i] = split_num[i]
    else:
        split_chars.reverse()
        split_num.reverse()
        for i in range(0, crosses):
            split_chars[i] = split_num[i]
        split_chars.reverse()
    new_segment = ''
    for element in split_chars:
        new_segment += element
    return new_segment

def min_num_to_fill(num_of_crosses, side):
    if num_of_crosses is 0 or side is 'r':
        return 0
    else:
        return int('1'+'0'*(num_of_crosses-1))

def max_num_to_fill(num_of_crosses):
    if num_of_crosses is 0:
        return 0
    return int('9'*num_of_crosses)

def is_empty(segment):
    """
    Проверяет, есть ли в сегменте цифры
    """
    for char in segment:
        if char != 'x':
            return False
    return True

def to_segment(str_seq, segment_len, shift):
    """
    Делит строку на list из значащих сегментов
    """
    ext_str_seq = segment_len*'x'+str_seq+segment_len*'x'
    segments = [ext_str_seq[num+shift:num+segment_len+shift]
                for num in range(0, len(ext_str_seq), segment_len)
                if not is_empty(ext_str_seq[num+shift:num+segment_len+shift])
                ]
    return segments

def get_rank_size(rank):
    if rank is 0:
        return 0
    size = ''
    for i in range(0, rank):
        size += '9'
    int_size = int(size)
    for i in range(1, rank):
        int_size -= get_rank_size(i)
    return int_size

def get_min_num(rank):
    min_num = '1'
    for i in range(1, rank):
        min_num += '0'
    return int(min_num)

def find_distance(num):
    summ = 0
    num_len = len(str(num))
    for rank in range(1, num_len):
        summ += rank * get_rank_size(rank)
    summ += (num-get_min_num(num_len))*num_len
    return summ

def are_equivalent(num, segment):
    if num == segment:
        return True
    digit_list = [char for char in segment]
    if len(digit_list) is 0:
        return True
    num_list = [char for char in num]
    if len(num_list) < len(digit_list):
        return False
    if segment[-1] is 'x':
        for i in range(0, len(digit_list)-1):
            if num_list[i] != digit_list[i] and digit_list[i] != 'x':
                return False
    else:
        num_list.reverse()
        digit_list.reverse()
        for i in range(0, len(digit_list)):
            if num_list[i] != digit_list[i] and digit_list[i] != 'x':
                return False
    return True

def min_num_in_segments(segments):
    if segments[0][0] is '0':
        return [float('Inf'), 0]
    if len(segments) is 1:
        if segments[0][0] is 'x':
            return [
                fill_crosses(segments[0],
                                str(min_num_to_fill(num_of_crosses(segments[0]), 'l'))),
                num_of_crosses(segments[0])]
        elif segments[0][-1] is 'x':
            return [
                fill_crosses(segments[0],
                                str(min_num_to_fill(num_of_crosses(segments[0]), 'r'))),
                num_of_crosses(segments[0])]
        return [segments[0],
                num_of_crosses(segments[0])]

    full_segments = [segment for segment in segments
                     if not 'x' in segment and len(segment) != 0]
    for segment in full_segments:
        if segment[0] is '0':
            return [float('Inf'), 0]
    if len(full_segments) > 1:
        for i in range(0, len(full_segments)-1):
            if int(full_segments[i])+1 != int(full_segments[i+1]):
                return [float('Inf'), 0]
    if len(full_segments) == len(segments):
        return [segments[0],
                num_of_crosses(segments[0])]

    if segments[0][0] is 'x':
        for i in range(min_num_to_fill(num_of_crosses(segments[0]), 'l'),
                       max_num_to_fill(num_of_crosses(segments[0]))):
            filled = fill_crosses(segments[0], str(i))
            correct = True
            for j in range(1, len(segments)):
                test = (str(int(filled)+j), segments[j])
                if not are_equivalent(str(int(filled)+j), segments[j]):
                    correct = False
            if correct is True:
                return [filled,
                        num_of_crosses(segments[0])]
    if segments[-1][-1] is 'x':
        for i in range(min_num_to_fill(num_of_crosses(segments[-1]), 'r'),
                       max_num_to_fill(num_of_crosses(segments[-1]))):
            filled = fill_crosses(segments[-1], str(i))
            correct = True
            for j in range(1, len(segments)):
                if not are_equivalent(str(int(filled)-j), segments[-j+1]):
                    correct = False
            if correct is True:
                return [filled,
                        num_of_crosses(segments[0])]
    return [float('Inf'), 0]

def split_seq(str_seq):
    nums = dict()
    for segment_len in range(1, len(str_seq)+1+num_of_zero(str_seq)):
        for shift in range(0, segment_len):
            segments = to_segment(str_seq, segment_len, shift)
            min_num = min_num_in_segments(segments)
            if min_num[0] != float('Inf'):
                nums[int(min_num[0])] = min_num[0]
            print(segments)
            print(min_num_in_segments(segments))
    return nums[min(nums.keys())]
#print(split_seq('125'))
# #str_seq = input('Введите искомую последовательность: ')
k = 1

while k < 10005:
    for i in range(k-1, k+2):
        print(i, end='  →  ')
        print(find_distance(i), end='  ~  ')
        print(first100.find(str(i)+str(i+1)+str(i+2)), end='  so  ')
        print('diff is: '+str(find_distance(i) -
                              first100.find(str(i)+str(i+1)+str(i+2))
                              ))
    k *= 10
<|endoftext|>


<commit_before>import os
import time
import pprint
import humanize
import functools

from datetime import datetime, timedelta
from holster.emitter import Priority
from disco.api.http import APIException
from disco.bot.command import CommandEvent, CommandLevels

from rowboat import BasePlugin as Plugin
from rowboat import RowboatPlugin, VERSION
from rowboat.sql import init_db
from rowboat.redis import rdb
from rowboat.models.guild import Guild
from rowboat.plugins.modlog import Actions

ENV = os.getenv('ENV', 'local')

PY_CODE_BLOCK = u'```py\n{}\n```'

INFO_MESSAGE = '''\
:information_source: Rowboat V{} - more information and detailed help can be found here:\
<https://github.com/b1naryth1ef/rowboat/wiki>
'''.format(VERSION)

class CorePlugin(Plugin):
    def load(self, ctx):
        init_db()

        self.startup = ctx.get('startup', datetime.utcnow())
        self.guilds = ctx.get('guilds', {})
        super(CorePlugin, self).load(ctx)

        for plugin in self.bot.plugins.values():
            if not isinstance(plugin, RowboatPlugin):
                continue

            plugin.register_trigger('command', 'pre', functools.partial(self.on_pre, plugin))
            plugin.register_trigger('listener', 'pre', functools.partial(self.on_pre, plugin))

    def unload(self, ctx):
        ctx['guilds'] = self.guilds
        ctx['startup'] = self.startup
        super(CorePlugin, self).unload(ctx)

    def on_pre(self, plugin, func, event, args, kwargs):
        if isinstance(event, CommandEvent):
            if event.command.metadata.get('global_', False):
                return event
        elif hasattr(func, 'subscriptions'):
            if func.subscriptions[0].metadata.get('global_', False):
                return event

        if hasattr(event, 'guild') and event.guild:
            guild_id = event.guild.id
        elif hasattr(event, 'guild_id') and event.guild_id:
            guild_id = event.guild_id
        else:
            return

        if guild_id not in self.guilds:
            return

        config = self.guilds[guild_id].get_config()

        plugin_name = plugin.name.lower().replace('plugin', '')
        if not getattr(config.plugins, plugin_name, None):
            return

        if plugin.whitelisted and plugin_name not in self.guilds[guild_id].whitelist:
            return

        event.config = getattr(config.plugins, plugin_name)
        return event

    @Plugin.schedule(290)
    def update_guild_bans(self):
        # TODO: disco calls this function right away, need guilds to be synced
        time.sleep(10)

        to_update = [
            guild for guild in Guild.select().where(
                (Guild.last_ban_sync < (datetime.utcnow() - timedelta(days=1))) |
                (Guild.last_ban_sync >> None)
            )
            if guild.guild_id in self.client.state.guilds]

        # Update 10 at a time
        for guild in to_update[:10]:
            guild.sync_bans(self.client.state.guilds.get(guild.guild_id))

    def send_control_message(self, content, *args, **kwargs):
        chan = self.client.state.dms.find_one(lambda dm: 80351110224678912 in dm.recipients)
        if not chan:
            self.log.warning('Failed to find control channel')
            return

        chan.send_message(u'[{}] {}'.format(ENV, content), *args, **kwargs)

    @Plugin.listen('Resumed')
    def on_resumed(self, event):
        self.send_control_message(u'Resumed: {}'.format(', '.join(event.trace)))

    @Plugin.listen('Ready')
    def on_ready(self, event):
        self.send_control_message(u'Connected: {}'.format(', '.join(event.trace)))

    @Plugin.listen('GuildCreate', priority=Priority.BEFORE, conditional=lambda e: not e.created)
    def on_guild_create(self, event):
        try:
            guild = Guild.with_id(event.id)
        except Guild.DoesNotExist:
            return

        # Ensure we're updated
        guild.sync(event.guild)

        self.guilds[event.id] = guild

        if guild.get_config().nickname:
            def set_nickname():
                m = event.members.select_one(id=self.state.me.id)
                if m and m.nick != guild.get_config().nickname:
                    try:
                        m.set_nickname(guild.get_config().nickname)
                    except APIException as e:
                        self.log.warning('Failed to set nickname for guild %s (%s)', event.guild, e.content)
            self.spawn_later(5, set_nickname)

    @Plugin.listen('MessageCreate')
    def on_message_create(self, event):
        # Ignore messages sent by bots
        if event.message.author.bot:
            return

        if rdb.sismember('ignored_channels', event.message.channel_id):
            return

        # If this is message for a guild, grab the guild object
        if hasattr(event, 'guild') and event.guild:
            guild_id = event.guild.id
        elif hasattr(event, 'guild_id') and event.guild_id:
            guild_id = event.guild_id
        else:
            guild_id = None

        guild = self.guilds.get(event.guild.id) if guild_id else None
        config = guild and guild.get_config()

        # If the guild has configuration, use that (otherwise use defaults)
        if config:
            if config.commands:
                commands = list(self.bot.get_commands_for_message(
                    config.commands.mention,
                    {},
                    config.commands.prefix,
                    event.message))
        elif guild_id:
            # Otherwise, default to requiring mentions
            commands = list(self.bot.get_commands_for_message(True, {}, '', event.message))
        else:
            if ENV != 'prod':
                if not event.message.content.startswith(ENV + '!'):
                    return
                event.message.content = event.message.content[len(ENV) + 1:]

            # DM's just use the commands (no prefix/mention)
            commands = list(self.bot.get_commands_for_message(False, {}, '', event.message))

        # If we didn't find any matching commands, return
        if not len(commands):
            return

        user_level = 0
        if config:
            for oid in event.guild.get_member(event.author).roles:
                if oid in config.levels and config.levels[oid] > user_level:
                    user_level = config.levels[oid]

            # User ID overrides should override all others
            if event.author.id in config.levels:
                user_level = config.levels[event.author.id]

        # Grab whether this user is a global admin
        # TODO: cache this
        global_admin = rdb.sismember('global_admins', event.author.id)

        # Iterate over commands and find a match
        for command, match in commands:
            if command.level == -1 and not global_admin:
                continue

            level = command.level

            if guild and not config and command.triggers[0] != 'setup':
                continue
            elif config and config.commands and command.plugin != self:
                if command.triggers[0] in config.commands.overrides:
                    override = config.commands.overrides[command.triggers[0]]
                    if override.disabled:
                        continue

                    if override.level is not None:
                        level = override.level

            if not global_admin and user_level < level:
                continue

            command.plugin.execute(CommandEvent(command, event.message, match))

            # Dispatch the command used modlog event
            if config:
                event.config = getattr(config.plugins, 'modlog', None)
                if not event.config:
                    return

                plugin = self.bot.plugins.get('ModLogPlugin')
                if plugin:
                    plugin.log_action(Actions.COMMAND_USED, event)

            return

    @Plugin.command('reload', '[plugin:str]', group='control', level=-1, oob=True)
    def command_control_reload(self, event, plugin=None):
        if not plugin:
            for plugin in self.bot.plugins.values():
                plugin.reload()
            return event.msg.reply(':recycle: reloaded all plugins')
        self.bot.plugins.get(plugin).reload()
        event.msg.reply(':recycle: reloaded plugin `{}`'.format(plugin))

    @Plugin.command('wl add', '<plugin:str> [guild:snowflake]', group='control', level=-1)
    def control_whitelist_add(self, event, plugin, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        guild.whitelist.append(plugin)
        guild.save()
        event.msg.reply(':ok_hand: this guild has been whitelisted for {}'.format(plugin))

    @Plugin.command('wl rmv', '<plugin:str> [guild:snowflake]', group='control', level=-1)
    def control_whitelist_rmv(self, event, plugin, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        if plugin not in guild.whitelist:
            return event.msg.reply(':warning: this guild isnt whitelisted for {}'.format(plugin))

        guild.whitelist.remove(plugin)
        guild.save()
        event.msg.reply(':ok_hand: this guild has been unwhitelisted for {}'.format(plugin))

    @Plugin.command('wl list', '[guild:snowflake]', group='control', level=-1)
    def control_whitelist_list(self, event, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        event.msg.reply('`{}`'.format(', '.join(guild.whitelist)))

    @Plugin.command('setup', '<url:str>')
    def command_setup(self, event, url):
        if not event.guild:
            return event.msg.reply(':warning: this command can only be used in servers')

        # Make sure we're not already setup
        if event.guild.id in self.guilds:
            return event.msg.reply(':warning: this server is already setup')

        global_admin = rdb.sismember('global_admins', event.author.id)

        # Make sure this is the owner of the server
        if not global_admin:
            if not event.guild.owner_id == event.author.id:
                return event.msg.reply(':warning: only the server owner can setup rowboat')

        # Make sure we have admin perms
        m = event.guild.members.select_one(id=self.state.me.id)
        if not m.permissions.administrator and not global_admin:
            return event.msg.reply(':warning: bot must have the Administrator permission')

        try:
            guild = Guild.create_from_url(event.guild, url)
            self.guilds[event.guild.id] = guild
            event.msg.reply(':ok_hand: successfully loaded configuration')
        except Exception as e:
            raise
            event.msg.reply(':no_entry: {}'.format(e))

    @Plugin.command('reload', level=CommandLevels.ADMIN)
    def command_reload(self, event):
        if not event.guild:
            return

        guild = self.guilds.get(event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild is not setup yet')

        try:
            guild.reload()
        except Exception as e:
            raise
            return event.msg.reply(':no_entry: {}'.format(e))

        event.msg.reply(':ok_hand: guild configuration reloaded')

    @Plugin.command('about', level=CommandLevels.ADMIN)
    def command_help(self, event):
        event.msg.reply(INFO_MESSAGE)

    @Plugin.command('config', level=CommandLevels.ADMIN)
    def command_config(self, event):
        if not event.guild or event.guild.id not in self.guilds:
            return

        event.msg.reply('Current configuration URL: <{}>'.format(self.guilds[event.guild.id].config_url))

    @Plugin.command('uptime', level=-1)
    def command_uptime(self, event):
        event.msg.reply('Rowboat was started {}'.format(
            humanize.naturaltime(datetime.utcnow() - self.startup)
        ))

    @Plugin.command('eval', level=-1)
    def command_eval(self, event):
        ctx = {
            'bot': self.bot,
            'client': self.bot.client,
            'state': self.bot.client.state,
            'event': event,
            'msg': event.msg,
            'guild': event.msg.guild,
            'channel': event.msg.channel,
            'author': event.msg.author
        }

        # Mulitline eval
        src = event.codeblock
        if src.count('\n'):
            lines = filter(bool, src.split('\n'))
            if lines[-1] and 'return' not in lines[-1]:
                lines[-1] = 'return ' + lines[-1]
            lines = '\n'.join('    ' + i for i in lines)
            code = 'def f():\n{}\nx = f()'.format(lines)
            local = {}

            try:
                exec compile(code, '<eval>', 'exec') in ctx, local
            except Exception as e:
                event.msg.reply(PY_CODE_BLOCK.format(type(e).__name__ + ': ' + str(e)))
                return

            event.msg.reply(PY_CODE_BLOCK.format(pprint.pformat(local['x'])))
        else:
            try:
                result = eval(src, ctx)
            except Exception as e:
                event.msg.reply(PY_CODE_BLOCK.format(type(e).__name__ + ': ' + str(e)))
                return

            event.msg.reply(PY_CODE_BLOCK.format(result))
<commit_msg>update command<commit_after>import os
import time
import pprint
import humanize
import functools
import subprocess

from datetime import datetime, timedelta
from holster.emitter import Priority
from disco.api.http import APIException
from disco.bot.command import CommandEvent, CommandLevels

from rowboat import BasePlugin as Plugin
from rowboat import RowboatPlugin, VERSION
from rowboat.sql import init_db
from rowboat.redis import rdb
from rowboat.models.guild import Guild
from rowboat.plugins.modlog import Actions

ENV = os.getenv('ENV', 'local')

PY_CODE_BLOCK = u'```py\n{}\n```'

INFO_MESSAGE = '''\
:information_source: Rowboat V{} - more information and detailed help can be found here:\
<https://github.com/b1naryth1ef/rowboat/wiki>
'''.format(VERSION)

class CorePlugin(Plugin):
    def load(self, ctx):
        init_db()

        self.startup = ctx.get('startup', datetime.utcnow())
        self.guilds = ctx.get('guilds', {})
        super(CorePlugin, self).load(ctx)

        for plugin in self.bot.plugins.values():
            if not isinstance(plugin, RowboatPlugin):
                continue

            plugin.register_trigger('command', 'pre', functools.partial(self.on_pre, plugin))
            plugin.register_trigger('listener', 'pre', functools.partial(self.on_pre, plugin))

    def unload(self, ctx):
        ctx['guilds'] = self.guilds
        ctx['startup'] = self.startup
        super(CorePlugin, self).unload(ctx)

    def on_pre(self, plugin, func, event, args, kwargs):
        if isinstance(event, CommandEvent):
            if event.command.metadata.get('global_', False):
                return event
        elif hasattr(func, 'subscriptions'):
            if func.subscriptions[0].metadata.get('global_', False):
                return event

        if hasattr(event, 'guild') and event.guild:
            guild_id = event.guild.id
        elif hasattr(event, 'guild_id') and event.guild_id:
            guild_id = event.guild_id
        else:
            return

        if guild_id not in self.guilds:
            return

        config = self.guilds[guild_id].get_config()

        plugin_name = plugin.name.lower().replace('plugin', '')
        if not getattr(config.plugins, plugin_name, None):
            return

        if plugin.whitelisted and plugin_name not in self.guilds[guild_id].whitelist:
            return

        event.config = getattr(config.plugins, plugin_name)
        return event

    @Plugin.schedule(290)
    def update_guild_bans(self):
        # TODO: disco calls this function right away, need guilds to be synced
        time.sleep(10)

        to_update = [
            guild for guild in Guild.select().where(
                (Guild.last_ban_sync < (datetime.utcnow() - timedelta(days=1))) |
                (Guild.last_ban_sync >> None)
            )
            if guild.guild_id in self.client.state.guilds]

        # Update 10 at a time
        for guild in to_update[:10]:
            guild.sync_bans(self.client.state.guilds.get(guild.guild_id))

    def send_control_message(self, content, *args, **kwargs):
        chan = self.client.state.dms.find_one(lambda dm: 80351110224678912 in dm.recipients)
        if not chan:
            self.log.warning('Failed to find control channel')
            return

        chan.send_message(u'[{}] {}'.format(ENV, content), *args, **kwargs)

    @Plugin.listen('Resumed')
    def on_resumed(self, event):
        self.send_control_message(u'Resumed: {}'.format(', '.join(event.trace)))

    @Plugin.listen('Ready')
    def on_ready(self, event):
        self.send_control_message(u'Connected: {}'.format(', '.join(event.trace)))

    @Plugin.listen('GuildCreate', priority=Priority.BEFORE, conditional=lambda e: not e.created)
    def on_guild_create(self, event):
        try:
            guild = Guild.with_id(event.id)
        except Guild.DoesNotExist:
            return

        # Ensure we're updated
        guild.sync(event.guild)

        self.guilds[event.id] = guild

        if guild.get_config().nickname:
            def set_nickname():
                m = event.members.select_one(id=self.state.me.id)
                if m and m.nick != guild.get_config().nickname:
                    try:
                        m.set_nickname(guild.get_config().nickname)
                    except APIException as e:
                        self.log.warning('Failed to set nickname for guild %s (%s)', event.guild, e.content)
            self.spawn_later(5, set_nickname)

    @Plugin.listen('MessageCreate')
    def on_message_create(self, event):
        # Ignore messages sent by bots
        if event.message.author.bot:
            return

        if rdb.sismember('ignored_channels', event.message.channel_id):
            return

        # If this is message for a guild, grab the guild object
        if hasattr(event, 'guild') and event.guild:
            guild_id = event.guild.id
        elif hasattr(event, 'guild_id') and event.guild_id:
            guild_id = event.guild_id
        else:
            guild_id = None

        guild = self.guilds.get(event.guild.id) if guild_id else None
        config = guild and guild.get_config()

        # If the guild has configuration, use that (otherwise use defaults)
        if config:
            if config.commands:
                commands = list(self.bot.get_commands_for_message(
                    config.commands.mention,
                    {},
                    config.commands.prefix,
                    event.message))
        elif guild_id:
            # Otherwise, default to requiring mentions
            commands = list(self.bot.get_commands_for_message(True, {}, '', event.message))
        else:
            if ENV != 'prod':
                if not event.message.content.startswith(ENV + '!'):
                    return
                event.message.content = event.message.content[len(ENV) + 1:]

            # DM's just use the commands (no prefix/mention)
            commands = list(self.bot.get_commands_for_message(False, {}, '', event.message))

        # If we didn't find any matching commands, return
        if not len(commands):
            return

        user_level = 0
        if config:
            for oid in event.guild.get_member(event.author).roles:
                if oid in config.levels and config.levels[oid] > user_level:
                    user_level = config.levels[oid]

            # User ID overrides should override all others
            if event.author.id in config.levels:
                user_level = config.levels[event.author.id]

        # Grab whether this user is a global admin
        # TODO: cache this
        global_admin = rdb.sismember('global_admins', event.author.id)

        # Iterate over commands and find a match
        for command, match in commands:
            if command.level == -1 and not global_admin:
                continue

            level = command.level

            if guild and not config and command.triggers[0] != 'setup':
                continue
            elif config and config.commands and command.plugin != self:
                if command.triggers[0] in config.commands.overrides:
                    override = config.commands.overrides[command.triggers[0]]
                    if override.disabled:
                        continue

                    if override.level is not None:
                        level = override.level

            if not global_admin and user_level < level:
                continue

            command.plugin.execute(CommandEvent(command, event.message, match))

            # Dispatch the command used modlog event
            if config:
                event.config = getattr(config.plugins, 'modlog', None)
                if not event.config:
                    return

                plugin = self.bot.plugins.get('ModLogPlugin')
                if plugin:
                    plugin.log_action(Actions.COMMAND_USED, event)

            return

    @Plugin.command('reload', '[plugin:str]', group='control', level=-1, oob=True)
    def command_control_reload(self, event, plugin=None):
        if not plugin:
            for plugin in self.bot.plugins.values():
                plugin.reload()
            return event.msg.reply(':recycle: reloaded all plugins')
        self.bot.plugins.get(plugin).reload()
        event.msg.reply(':recycle: reloaded plugin `{}`'.format(plugin))

    @Plugin.command('update', group='control', level=-1, oob=True)
    def commnad_update(self, event):
        proc = subprocess.Popen(['git', 'pull', 'origin', 'master'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)

        if proc.returncode != 0:
            return event.msg.reply(':warning: failed to update: ```{}```'.format(proc.stderr.read()))

        for plugin in self.bot.plugins.values():
            plugin.reload()
        event.msg.reply(':ok_hand: updated: ```{}```'.format(proc.stdout.read()))

    @Plugin.command('wl add', '<plugin:str> [guild:snowflake]', group='control', level=-1)
    def control_whitelist_add(self, event, plugin, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        guild.whitelist.append(plugin)
        guild.save()
        event.msg.reply(':ok_hand: this guild has been whitelisted for {}'.format(plugin))

    @Plugin.command('wl rmv', '<plugin:str> [guild:snowflake]', group='control', level=-1)
    def control_whitelist_rmv(self, event, plugin, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        if plugin not in guild.whitelist:
            return event.msg.reply(':warning: this guild isnt whitelisted for {}'.format(plugin))

        guild.whitelist.remove(plugin)
        guild.save()
        event.msg.reply(':ok_hand: this guild has been unwhitelisted for {}'.format(plugin))

    @Plugin.command('wl list', '[guild:snowflake]', group='control', level=-1)
    def control_whitelist_list(self, event, guild=None):
        guild = self.guilds.get(guild or event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild isnt setup yet')

        event.msg.reply('`{}`'.format(', '.join(guild.whitelist)))

    @Plugin.command('setup', '<url:str>')
    def command_setup(self, event, url):
        if not event.guild:
            return event.msg.reply(':warning: this command can only be used in servers')

        # Make sure we're not already setup
        if event.guild.id in self.guilds:
            return event.msg.reply(':warning: this server is already setup')

        global_admin = rdb.sismember('global_admins', event.author.id)

        # Make sure this is the owner of the server
        if not global_admin:
            if not event.guild.owner_id == event.author.id:
                return event.msg.reply(':warning: only the server owner can setup rowboat')

        # Make sure we have admin perms
        m = event.guild.members.select_one(id=self.state.me.id)
        if not m.permissions.administrator and not global_admin:
            return event.msg.reply(':warning: bot must have the Administrator permission')

        try:
            guild = Guild.create_from_url(event.guild, url)
            self.guilds[event.guild.id] = guild
            event.msg.reply(':ok_hand: successfully loaded configuration')
        except Exception as e:
            raise
            event.msg.reply(':no_entry: {}'.format(e))

    @Plugin.command('reload', level=CommandLevels.ADMIN)
    def command_reload(self, event):
        if not event.guild:
            return

        guild = self.guilds.get(event.guild.id)
        if not guild:
            return event.msg.reply(':warning: this guild is not setup yet')

        try:
            guild.reload()
        except Exception as e:
            raise
            return event.msg.reply(':no_entry: {}'.format(e))

        event.msg.reply(':ok_hand: guild configuration reloaded')

    @Plugin.command('about', level=CommandLevels.ADMIN)
    def command_help(self, event):
        event.msg.reply(INFO_MESSAGE)

    @Plugin.command('config', level=CommandLevels.ADMIN)
    def command_config(self, event):
        if not event.guild or event.guild.id not in self.guilds:
            return

        event.msg.reply('Current configuration URL: <{}>'.format(self.guilds[event.guild.id].config_url))

    @Plugin.command('uptime', level=-1)
    def command_uptime(self, event):
        event.msg.reply('Rowboat was started {}'.format(
            humanize.naturaltime(datetime.utcnow() - self.startup)
        ))

    @Plugin.command('eval', level=-1)
    def command_eval(self, event):
        ctx = {
            'bot': self.bot,
            'client': self.bot.client,
            'state': self.bot.client.state,
            'event': event,
            'msg': event.msg,
            'guild': event.msg.guild,
            'channel': event.msg.channel,
            'author': event.msg.author
        }

        # Mulitline eval
        src = event.codeblock
        if src.count('\n'):
            lines = filter(bool, src.split('\n'))
            if lines[-1] and 'return' not in lines[-1]:
                lines[-1] = 'return ' + lines[-1]
            lines = '\n'.join('    ' + i for i in lines)
            code = 'def f():\n{}\nx = f()'.format(lines)
            local = {}

            try:
                exec compile(code, '<eval>', 'exec') in ctx, local
            except Exception as e:
                event.msg.reply(PY_CODE_BLOCK.format(type(e).__name__ + ': ' + str(e)))
                return

            event.msg.reply(PY_CODE_BLOCK.format(pprint.pformat(local['x'])))
        else:
            try:
                result = eval(src, ctx)
            except Exception as e:
                event.msg.reply(PY_CODE_BLOCK.format(type(e).__name__ + ': ' + str(e)))
                return

            event.msg.reply(PY_CODE_BLOCK.format(result))
<|endoftext|>