<commit_before># Copyright (c) 2010-2012 OpenStack, LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cPickle as pickle
import os
import signal
import sys
import time
from random import random

from eventlet import patcher, Timeout

from swift.common.bufferedhttp import http_connect
from swift.common.exceptions import ConnectionTimeout
from swift.common.ring import Ring
from swift.common.utils import get_logger, renamer, write_pickle
from swift.common.daemon import Daemon
from swift.obj.server import ASYNCDIR

class ObjectUpdater(Daemon):
    """Update object information in container listings."""

    def __init__(self, conf):
        self.conf = conf
        self.logger = get_logger(conf, log_route='object-updater')
        self.devices = conf.get('devices', '/srv/node')
        self.mount_check = conf.get('mount_check', 'true').lower() in \
                              ('true', 't', '1', 'on', 'yes', 'y')
        swift_dir = conf.get('swift_dir', '/etc/swift')
        self.interval = int(conf.get('interval', 300))
        self.container_ring_path = os.path.join(swift_dir, 'container.ring.gz')
        self.container_ring = None
        self.concurrency = int(conf.get('concurrency', 1))
        self.slowdown = float(conf.get('slowdown', 0.01))
        self.node_timeout = int(conf.get('node_timeout', 10))
        self.conn_timeout = float(conf.get('conn_timeout', 0.5))
        self.successes = 0
        self.failures = 0

    def get_container_ring(self):
        """Get the container ring.  Load it, if it hasn't been yet."""
        if not self.container_ring:
            self.logger.debug(
                _('Loading container ring from %s'), self.container_ring_path)
            self.container_ring = Ring(self.container_ring_path)
        return self.container_ring

    def run_forever(self, *args, **kwargs):
        """Run the updater continuously."""
        time.sleep(random() * self.interval)
        while True:
            self.logger.info(_('Begin object update sweep'))
            begin = time.time()
            pids = []
            # read from container ring to ensure it's fresh
            self.get_container_ring().get_nodes('')
            for device in os.listdir(self.devices):
                if self.mount_check and not \
                        os.path.ismount(os.path.join(self.devices, device)):
                    self.logger.warn(
                        _('Skipping %s as it is not mounted'), device)
                    continue
                while len(pids) >= self.concurrency:
                    pids.remove(os.wait()[0])
                pid = os.fork()
                if pid:
                    pids.append(pid)
                else:
                    signal.signal(signal.SIGTERM, signal.SIG_DFL)
                    patcher.monkey_patch(all=False, socket=True)
                    self.successes = 0
                    self.failures = 0
                    forkbegin = time.time()
                    self.object_sweep(os.path.join(self.devices, device))
                    elapsed = time.time() - forkbegin
                    self.logger.info(_('Object update sweep of %(device)s'
                        ' completed: %(elapsed).02fs, %(success)s successes'
                        ', %(fail)s failures'),
                        {'device': device, 'elapsed': elapsed,
                         'success': self.successes, 'fail': self.failures})
                    sys.exit()
            while pids:
                pids.remove(os.wait()[0])
            elapsed = time.time() - begin
            self.logger.info(_('Object update sweep completed: %.02fs'),
                    elapsed)
            if elapsed < self.interval:
                time.sleep(self.interval - elapsed)

    def run_once(self, *args, **kwargs):
        """Run the updater once"""
        self.logger.info(_('Begin object update single threaded sweep'))
        begin = time.time()
        self.successes = 0
        self.failures = 0
        for device in os.listdir(self.devices):
            if self.mount_check and \
                    not os.path.ismount(os.path.join(self.devices, device)):
                self.logger.warn(
                    _('Skipping %s as it is not mounted'), device)
                continue
            self.object_sweep(os.path.join(self.devices, device))
        elapsed = time.time() - begin
        self.logger.info(_('Object update single threaded sweep completed: '
            '%(elapsed).02fs, %(success)s successes, %(fail)s failures'),
            {'elapsed': elapsed, 'success': self.successes,
             'fail': self.failures})

    def object_sweep(self, device):
        """
        If there are async pendings on the device, walk each one and update.

        :param device: path to device
        """
        async_pending = os.path.join(device, ASYNCDIR)
        if not os.path.isdir(async_pending):
            return
        for prefix in os.listdir(async_pending):
            prefix_path = os.path.join(async_pending, prefix)
            if not os.path.isdir(prefix_path):
                continue
            last_obj_hash = None
            for update in sorted(os.listdir(prefix_path), reverse=True):
                update_path = os.path.join(prefix_path, update)
                if not os.path.isfile(update_path):
                    continue
                try:
                    obj_hash, timestamp = update.split('-')
                except ValueError:
                    self.logger.error(
                        _('ERROR async pending file with unexpected name %s')
                        % (update_path))
                    continue
                if obj_hash == last_obj_hash:
                    os.unlink(update_path)
                else:
                    self.process_object_update(update_path, device)
                    last_obj_hash = obj_hash
                time.sleep(self.slowdown)
            try:
                os.rmdir(prefix_path)
            except OSError:
                pass

    def process_object_update(self, update_path, device):
        """
        Process the object information to be updated and update.

        :param update_path: path to pickled object update file
        :param device: path to device
        """
        try:
            update = pickle.load(open(update_path, 'rb'))
        except Exception:
            self.logger.exception(
                _('ERROR Pickle problem, quarantining %s'), update_path)
            renamer(update_path, os.path.join(device,
                'quarantined', 'objects', os.path.basename(update_path)))
            return
        successes = update.get('successes', [])
        part, nodes = self.get_container_ring().get_nodes(
                                update['account'], update['container'])
        obj = '/%s/%s/%s' % \
              (update['account'], update['container'], update['obj'])
        success = True
        for node in nodes:
            if node['id'] not in successes:
                status = self.object_update(node, part, update['op'], obj,
                                        update['headers'])
                if not (200 <= status < 300) and status != 404:
                    success = False
                else:
                    successes.append(node['id'])
        if success:
            self.successes += 1
            self.logger.debug(_('Update sent for %(obj)s %(path)s'),
                {'obj': obj, 'path': update_path})
            os.unlink(update_path)
        else:
            self.failures += 1
            self.logger.debug(_('Update failed for %(obj)s %(path)s'),
                {'obj': obj, 'path': update_path})
            update['successes'] = successes
            write_pickle(update, update_path, os.path.join(device, 'tmp'))

    def object_update(self, node, part, op, obj, headers):
        """
        Perform the object update to the container

        :param node: node dictionary from the container ring
        :param part: partition that holds the container
        :param op: operation performed (ex: 'POST' or 'DELETE')
        :param obj: object name being updated
        :param headers: headers to send with the update
        """
        try:
            with ConnectionTimeout(self.conn_timeout):
                conn = http_connect(node['ip'], node['port'], node['device'],
                    part, op, obj, headers)
            with Timeout(self.node_timeout):
                resp = conn.getresponse()
                resp.read()
                return resp.status
        except (Exception, Timeout):
            self.logger.exception(_('ERROR with remote server '
                '%(ip)s:%(port)s/%(device)s'), node)
        return 500
<commit_msg>Only rewrite async-pending files if the data changed<commit_after># Copyright (c) 2010-2012 OpenStack, LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cPickle as pickle
import os
import signal
import sys
import time
from random import random

from eventlet import patcher, Timeout

from swift.common.bufferedhttp import http_connect
from swift.common.exceptions import ConnectionTimeout
from swift.common.ring import Ring
from swift.common.utils import get_logger, renamer, write_pickle
from swift.common.daemon import Daemon
from swift.obj.server import ASYNCDIR

class ObjectUpdater(Daemon):
    """Update object information in container listings."""

    def __init__(self, conf):
        self.conf = conf
        self.logger = get_logger(conf, log_route='object-updater')
        self.devices = conf.get('devices', '/srv/node')
        self.mount_check = conf.get('mount_check', 'true').lower() in \
                              ('true', 't', '1', 'on', 'yes', 'y')
        swift_dir = conf.get('swift_dir', '/etc/swift')
        self.interval = int(conf.get('interval', 300))
        self.container_ring_path = os.path.join(swift_dir, 'container.ring.gz')
        self.container_ring = None
        self.concurrency = int(conf.get('concurrency', 1))
        self.slowdown = float(conf.get('slowdown', 0.01))
        self.node_timeout = int(conf.get('node_timeout', 10))
        self.conn_timeout = float(conf.get('conn_timeout', 0.5))
        self.successes = 0
        self.failures = 0

    def get_container_ring(self):
        """Get the container ring.  Load it, if it hasn't been yet."""
        if not self.container_ring:
            self.logger.debug(
                _('Loading container ring from %s'), self.container_ring_path)
            self.container_ring = Ring(self.container_ring_path)
        return self.container_ring

    def run_forever(self, *args, **kwargs):
        """Run the updater continuously."""
        time.sleep(random() * self.interval)
        while True:
            self.logger.info(_('Begin object update sweep'))
            begin = time.time()
            pids = []
            # read from container ring to ensure it's fresh
            self.get_container_ring().get_nodes('')
            for device in os.listdir(self.devices):
                if self.mount_check and not \
                        os.path.ismount(os.path.join(self.devices, device)):
                    self.logger.warn(
                        _('Skipping %s as it is not mounted'), device)
                    continue
                while len(pids) >= self.concurrency:
                    pids.remove(os.wait()[0])
                pid = os.fork()
                if pid:
                    pids.append(pid)
                else:
                    signal.signal(signal.SIGTERM, signal.SIG_DFL)
                    patcher.monkey_patch(all=False, socket=True)
                    self.successes = 0
                    self.failures = 0
                    forkbegin = time.time()
                    self.object_sweep(os.path.join(self.devices, device))
                    elapsed = time.time() - forkbegin
                    self.logger.info(_('Object update sweep of %(device)s'
                        ' completed: %(elapsed).02fs, %(success)s successes'
                        ', %(fail)s failures'),
                        {'device': device, 'elapsed': elapsed,
                         'success': self.successes, 'fail': self.failures})
                    sys.exit()
            while pids:
                pids.remove(os.wait()[0])
            elapsed = time.time() - begin
            self.logger.info(_('Object update sweep completed: %.02fs'),
                    elapsed)
            if elapsed < self.interval:
                time.sleep(self.interval - elapsed)

    def run_once(self, *args, **kwargs):
        """Run the updater once"""
        self.logger.info(_('Begin object update single threaded sweep'))
        begin = time.time()
        self.successes = 0
        self.failures = 0
        for device in os.listdir(self.devices):
            if self.mount_check and \
                    not os.path.ismount(os.path.join(self.devices, device)):
                self.logger.warn(
                    _('Skipping %s as it is not mounted'), device)
                continue
            self.object_sweep(os.path.join(self.devices, device))
        elapsed = time.time() - begin
        self.logger.info(_('Object update single threaded sweep completed: '
            '%(elapsed).02fs, %(success)s successes, %(fail)s failures'),
            {'elapsed': elapsed, 'success': self.successes,
             'fail': self.failures})

    def object_sweep(self, device):
        """
        If there are async pendings on the device, walk each one and update.

        :param device: path to device
        """
        async_pending = os.path.join(device, ASYNCDIR)
        if not os.path.isdir(async_pending):
            return
        for prefix in os.listdir(async_pending):
            prefix_path = os.path.join(async_pending, prefix)
            if not os.path.isdir(prefix_path):
                continue
            last_obj_hash = None
            for update in sorted(os.listdir(prefix_path), reverse=True):
                update_path = os.path.join(prefix_path, update)
                if not os.path.isfile(update_path):
                    continue
                try:
                    obj_hash, timestamp = update.split('-')
                except ValueError:
                    self.logger.error(
                        _('ERROR async pending file with unexpected name %s')
                        % (update_path))
                    continue
                if obj_hash == last_obj_hash:
                    os.unlink(update_path)
                else:
                    self.process_object_update(update_path, device)
                    last_obj_hash = obj_hash
                time.sleep(self.slowdown)
            try:
                os.rmdir(prefix_path)
            except OSError:
                pass

    def process_object_update(self, update_path, device):
        """
        Process the object information to be updated and update.

        :param update_path: path to pickled object update file
        :param device: path to device
        """
        try:
            update = pickle.load(open(update_path, 'rb'))
        except Exception:
            self.logger.exception(
                _('ERROR Pickle problem, quarantining %s'), update_path)
            renamer(update_path, os.path.join(device,
                'quarantined', 'objects', os.path.basename(update_path)))
            return
        successes = update.get('successes', [])
        part, nodes = self.get_container_ring().get_nodes(
                                update['account'], update['container'])
        obj = '/%s/%s/%s' % \
              (update['account'], update['container'], update['obj'])
        success = True
        new_successes = False
        for node in nodes:
            if node['id'] not in successes:
                status = self.object_update(node, part, update['op'], obj,
                                        update['headers'])
                if not (200 <= status < 300) and status != 404:
                    success = False
                else:
                    successes.append(node['id'])
                    new_successes = True
        if success:
            self.successes += 1
            self.logger.debug(_('Update sent for %(obj)s %(path)s'),
                {'obj': obj, 'path': update_path})
            os.unlink(update_path)
        else:
            self.failures += 1
            self.logger.debug(_('Update failed for %(obj)s %(path)s'),
                {'obj': obj, 'path': update_path})
            if new_successes:
                update['successes'] = successes
                write_pickle(update, update_path, os.path.join(device, 'tmp'))

    def object_update(self, node, part, op, obj, headers):
        """
        Perform the object update to the container

        :param node: node dictionary from the container ring
        :param part: partition that holds the container
        :param op: operation performed (ex: 'POST' or 'DELETE')
        :param obj: object name being updated
        :param headers: headers to send with the update
        """
        try:
            with ConnectionTimeout(self.conn_timeout):
                conn = http_connect(node['ip'], node['port'], node['device'],
                    part, op, obj, headers)
            with Timeout(self.node_timeout):
                resp = conn.getresponse()
                resp.read()
                return resp.status
        except (Exception, Timeout):
            self.logger.exception(_('ERROR with remote server '
                '%(ip)s:%(port)s/%(device)s'), node)
        return 500
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import gc
import riak

from bottle import Bottle, request, response
from bottle.ext.mongo import MongoPlugin

from pandas import DataFrame

from settings import PROJECT_PATH
from utils import df_generate, conf

export_app = Bottle()
mongo = MongoPlugin(
    uri=conf("mongodb")["uri"],
    db=conf("mongodb")["db"],
    json_mongo=True)
export_app.install(mongo)

@export_app.route('/data/<slug>.<ext>')
def data(mongodb, slug, ext='xls'):
    MyClient = riak.RiakClient(protocol=conf("riak")["protocol"],
                               http_port=conf("riak")["http_port"],
                               host=conf("riak")["host"])

    MyBucket = MyClient.bucket(conf("riak")["bucket"])

    element = mongodb['element'].find_one({'slug': slug})

    columns = json.loads(MyBucket.get(
        '{}-columns'.format(element.get('cube'))).data or [])

    fields = columns
    if request.GET.get('fields', None):
        fields = request.GET.get('fields').split(',')

    filters = [i[0] for i in request.GET.iteritems()
               if len(i[0].split('filter__')) > 1]

    df = DataFrame(MyBucket.get(element.get('cube')).data, columns=fields)
    if len(filters) >= 1:
        for f in filters:
            df = df.query(df_generate(df, request.GET.get(f), f))

    groupby = []
    if request.GET.get('groupby', None):
        groupby = request.GET.get('groupby', ).split(',')
    if len(groupby) >= 1:
        df = df.groupby(groupby)

    if request.GET.get('orderby', None):
        orderby = request.GET.get('orderby', [])
        orderby__order = True
        if request.GET.get('orderby__order', 0) != 1:
            orderby__order = False
        df = df.sort(orderby, ascending=orderby__order)

    # CLEAN MEMORY
    del filters, fields, columns
    gc.collect()

    file_name = '{}/assets/exports/openmining-{}.{}'.format(
        PROJECT_PATH, element.get('cube'), ext)
    if ext == 'csv':
        df.to_csv(file_name)
        contenttype = 'text/csv'
    else:
        df.to_excel(file_name)
        contenttype = 'application/vnd.ms-excel'

    response.set_header('charset', 'utf-8')
    response.set_header('Content-disposition', 'attachment; '
                        'filename={}.{}'.format(element.get('cube'), ext))
    response.content_type = contenttype

    ifile = open(file_name, "r")
    o = ifile.read()
    ifile.close()
    return o
<commit_msg>Used DataFrameSearchColumn on export api<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import gc
import riak

from bottle import Bottle, request, response
from bottle.ext.mongo import MongoPlugin

from pandas import DataFrame

from settings import PROJECT_PATH
from utils import df_generate, conf, DataFrameSearchColumn

export_app = Bottle()
mongo = MongoPlugin(
    uri=conf("mongodb")["uri"],
    db=conf("mongodb")["db"],
    json_mongo=True)
export_app.install(mongo)

@export_app.route('/data/<slug>.<ext>')
def data(mongodb, slug, ext='xls'):
    MyClient = riak.RiakClient(protocol=conf("riak")["protocol"],
                               http_port=conf("riak")["http_port"],
                               host=conf("riak")["host"])

    MyBucket = MyClient.bucket(conf("riak")["bucket"])

    element = mongodb['element'].find_one({'slug': slug})

    columns = json.loads(MyBucket.get(
        '{}-columns'.format(element.get('cube'))).data or [])

    fields = columns
    if request.GET.get('fields', None):
        fields = request.GET.get('fields').split(',')

    filters = [i[0] for i in request.GET.iteritems()
               if len(i[0].split('filter__')) > 1]

    df = DataFrame(MyBucket.get(element.get('cube')).data, columns=fields)
    if len(filters) >= 1:
        for f in filters:
            s = f.split('__')
            field = s[1]
            operator = s[2]
            value = request.GET.get(f)
            if operator == 'like':
                df = DataFrameSearchColumn(df, field, value)
            else:
                df = df.query(df_generate(df, value, f))

    groupby = []
    if request.GET.get('groupby', None):
        groupby = request.GET.get('groupby', ).split(',')
    if len(groupby) >= 1:
        df = df.groupby(groupby)

    if request.GET.get('orderby', None):
        orderby = request.GET.get('orderby', [])
        orderby__order = True
        if request.GET.get('orderby__order', 0) != 1:
            orderby__order = False
        df = df.sort(orderby, ascending=orderby__order)

    # CLEAN MEMORY
    del filters, fields, columns
    gc.collect()

    file_name = '{}/assets/exports/openmining-{}.{}'.format(
        PROJECT_PATH, element.get('cube'), ext)
    if ext == 'csv':
        df.to_csv(file_name)
        contenttype = 'text/csv'
    else:
        df.to_excel(file_name)
        contenttype = 'application/vnd.ms-excel'

    response.set_header('charset', 'utf-8')
    response.set_header('Content-disposition', 'attachment; '
                        'filename={}.{}'.format(element.get('cube'), ext))
    response.content_type = contenttype

    ifile = open(file_name, "r")
    o = ifile.read()
    ifile.close()
    return o
<|endoftext|>


<commit_before># $Id$
#
# PyGEOS distutils setup script
#
# Copyright 2004 Sean Gillies, sgillies@frii.com
#
# This is free software; you can redistribute and/or modify it under
# the terms of the GNU Lesser General Public Licence as published
# by the Free Software Foundation.
# See the COPYING file for more information.
#
# ===========================================================================
#
# Test before installation
#
# $ python setup.py build
# $ python tests/runtests.py -v
# $ sudo python setup.py install
#
# ===========================================================================

from distutils.core import setup, Extension

# We're skipping the metadata for now

# the extension module, source was generated by SWIG from geos.i
#
# Note that until geos is ready to be installed, we're linking against
# the pre-installed library in source/geom/.libs.
extmod = Extension("_geos", ["geos_wrap.cxx"],
                   include_dirs = ["../../source/headers",
                                   "../../source/headers/geos"],
                   libraries = ["stdc++"],
                   extra_objects = ["../../source/geom/.libs/libgeos.a"])

# the python module, generated by SWIG from geos.i
pymod = "geos"

# The setup function
setup(name = "PyGEOS", ext_modules = [extmod], py_modules = [pymod])

<commit_msg>now supporting the win32 platform, thanks to Howard Butler<commit_after># $Id$
#
# PyGEOS distutils setup script
#
# Copyright 2004 Sean Gillies, sgillies@frii.com
#
# This is free software; you can redistribute and/or modify it under
# the terms of the GNU Lesser General Public Licence as published
# by the Free Software Foundation.
# See the COPYING file for more information.
#
# ===========================================================================
#
# Test before installation
#
# $ python setup.py build
# $ python tests/runtests.py -v
# $ sudo python setup.py install
#
# ===========================================================================

from distutils.core import setup, Extension
import sys

# the extension module, source was generated by SWIG from geos.i
#
# Note that until geos is ready to be installed, we're linking against
# the pre-installed library in source/geom/.libs.

if sys.platform=='win32':
    extmod = Extension('_geos', ['geos_wrap.cxx'],
                       include_dirs = ['../../source/headers',
                                       '../../source/headers/geos'],
                       libraries = ['geos'],
                       library_dirs=['../../source'],)

else:
    extmod = Extension('_geos', ['geos_wrap.cxx'],
                       include_dirs = ['../../source/headers',
                                       '../../source/headers/geos'],
                       libraries = ['stdc++'],
                       extra_objects = ['../../source/geom/.libs/libgeos.a'])

# the python module, generated by SWIG from geos.i
pymod = 'geos'

# The setup function
setup(name = 'PyGEOS',
      version = '2.1.1',
      description = 'Python SWIG interface for GEOS',
      url='http://geos.refractions.net',
      ext_modules = [extmod],
      py_modules = [pymod])

<|endoftext|>


<commit_before>from attendee_tournaments import *

AutomatedEmail.queries[AttendeeTournament] = lambda session: session.query(AttendeeTournament).all()

AutomatedEmail(AttendeeTournament, 'Your tournament application has been received', 'tournament_app_received.txt',
               lambda app: app.status == c.NEW)

AutomatedEmail(AttendeeTournament, 'Your tournament application has been approved', 'tournament_app_accepted.txt',
               lambda app: app.status == c.ACCEPTED)

AutomatedEmail(AttendeeTournament, 'Your tournament application has been declined', 'tournament_app_declined.txt',
               lambda app: app.status == c.DECLINED)
<commit_msg>Fixes #2383: Adds static ident arg to each AutomatedEmail instance<commit_after>from attendee_tournaments import *

AutomatedEmail.queries[AttendeeTournament] = lambda session: session.query(AttendeeTournament).all()

AutomatedEmail(AttendeeTournament, 'Your tournament application has been received', 'tournament_app_received.txt',
               lambda app: app.status == c.NEW,
               ident='tournament_app_received')

AutomatedEmail(AttendeeTournament, 'Your tournament application has been approved', 'tournament_app_accepted.txt',
               lambda app: app.status == c.ACCEPTED,
               ident='tournament_app_accepted')

AutomatedEmail(AttendeeTournament, 'Your tournament application has been declined', 'tournament_app_declined.txt',
               lambda app: app.status == c.DECLINED,
               ident='tournament_app_declined')
<|endoftext|>


<commit_before>from PyQt5.QtWidgets import QDialog, QTabWidget, QVBoxLayout
from PyQt5.QtCore import QObject, pyqtSignal
from sakia.decorators import asyncify
from sakia.gui.widgets.dialogs import dialog_async_exec, QAsyncMessageBox
from .model import UserInformationModel
from .view import UserInformationView
import logging

class UserInformationController(QObject):
    """
    The homescreen view
    """
    identity_loaded = pyqtSignal()

    def __init__(self, parent, view, model):
        """
        Constructor of the homescreen component

        :param sakia.gui.homescreen.view.HomeScreenView: the view
        :param sakia.gui.homescreen.model.HomeScreenModel model: the model
        """
        super().__init__(parent)
        self.view = view
        self.model = model
        self._logger = logging.getLogger('sakia')

    @classmethod
    def create(cls, parent, app, identity):
        view = UserInformationView(parent.view if parent else None)
        model = UserInformationModel(None, app, identity)
        homescreen = cls(parent, view, model)
        model.setParent(homescreen)
        return homescreen

    @classmethod
    def show_identity(cls, parent, app, identity):
        dialog = QDialog()
        dialog.setWindowTitle("Informations")
        user_info = cls.create(parent, app, identity)
        user_info.view.setParent(dialog)
        dialog.setLayout(QVBoxLayout(dialog))
        dialog.layout().addWidget(user_info.view)
        user_info.refresh()
        dialog.exec()

    @classmethod
    @asyncify
    async def search_and_show_pubkey(cls, parent, app, pubkey):
        dialog = QDialog(parent)
        dialog.setWindowTitle("Informations")
        layout = QVBoxLayout(dialog)
        tabwidget = QTabWidget(dialog)
        layout.addWidget(tabwidget)

        identities = await app.identities_service.lookup(pubkey)
        for i in identities:
            user_info = cls.create(parent, app, i)
            user_info.refresh()
            tabwidget.addTab(user_info.view, i.uid)
        return await dialog_async_exec(dialog)

    @asyncify
    async def refresh(self):
        try:
            if self.model.identity:
                self.view.show_busy()
                await self.model.load_identity(self.model.identity)
                self.view.display_uid(self.model.identity.uid, self.model.identity.member)
                self.view.display_identity_timestamps(self.model.identity.pubkey, self.model.identity.timestamp,
                                                      self.model.identity.membership_timestamp,
                                                      self.model.mstime_remaining(), await self.model.nb_certs())
                self.view.hide_busy()
                self.identity_loaded.emit()
        except RuntimeError as e:
            # object can be deleted by Qt during asynchronous ops
            # we don't care of this error
            if "wrapped C/C++ object" in str(e):
                self._logger.debug(str(e))
                pass
            else:
                raise

    @asyncify
    async def search_identity(self, identity):
        await self.model.load_identity(identity)
        self.refresh()

    def change_identity(self, identity):
        """
        Set identity
        :param sakia.core.registry.Identity identity:
        """
        self.model.identity = identity
        self.refresh()
<commit_msg>Fix user information dialog size<commit_after>from PyQt5.QtWidgets import QDialog, QTabWidget, QVBoxLayout
from PyQt5.QtCore import QObject, pyqtSignal
from sakia.decorators import asyncify
from sakia.gui.widgets.dialogs import dialog_async_exec, QAsyncMessageBox
from .model import UserInformationModel
from .view import UserInformationView
import logging

class UserInformationController(QObject):
    """
    The homescreen view
    """
    identity_loaded = pyqtSignal()

    def __init__(self, parent, view, model):
        """
        Constructor of the homescreen component

        :param sakia.gui.homescreen.view.HomeScreenView: the view
        :param sakia.gui.homescreen.model.HomeScreenModel model: the model
        """
        super().__init__(parent)
        self.view = view
        self.model = model
        self._logger = logging.getLogger('sakia')

    @classmethod
    def create(cls, parent, app, identity):
        view = UserInformationView(parent.view if parent else None)
        model = UserInformationModel(None, app, identity)
        homescreen = cls(parent, view, model)
        model.setParent(homescreen)
        return homescreen

    @classmethod
    def show_identity(cls, parent, app, identity):
        dialog = QDialog()
        dialog.setWindowTitle("Informations")
        user_info = cls.create(parent, app, identity)
        user_info.view.setParent(dialog)
        dialog.setLayout(QVBoxLayout(dialog))
        dialog.layout().addWidget(user_info.view)
        user_info.refresh()
        dialog.resize(800, 300)
        dialog.exec()

    @classmethod
    @asyncify
    async def search_and_show_pubkey(cls, parent, app, pubkey):
        dialog = QDialog(parent)
        dialog.setWindowTitle("Informations")
        layout = QVBoxLayout(dialog)
        tabwidget = QTabWidget(dialog)
        layout.addWidget(tabwidget)

        identities = await app.identities_service.lookup(pubkey)
        for i in identities:
            user_info = cls.create(parent, app, i)
            user_info.refresh()
            tabwidget.addTab(user_info.view, i.uid)
        return await dialog_async_exec(dialog)

    @asyncify
    async def refresh(self):
        try:
            if self.model.identity:
                self.view.show_busy()
                await self.model.load_identity(self.model.identity)
                self.view.display_uid(self.model.identity.uid, self.model.identity.member)
                self.view.display_identity_timestamps(self.model.identity.pubkey, self.model.identity.timestamp,
                                                      self.model.identity.membership_timestamp,
                                                      self.model.mstime_remaining(), await self.model.nb_certs())
                self.view.hide_busy()
                self.identity_loaded.emit()
        except RuntimeError as e:
            # object can be deleted by Qt during asynchronous ops
            # we don't care of this error
            if "wrapped C/C++ object" in str(e):
                self._logger.debug(str(e))
                pass
            else:
                raise

    @asyncify
    async def search_identity(self, identity):
        await self.model.load_identity(identity)
        self.refresh()

    def change_identity(self, identity):
        """
        Set identity
        :param sakia.core.registry.Identity identity:
        """
        self.model.identity = identity
        self.refresh()
<|endoftext|>


<commit_before>import datetime
import re
import urllib2
import urlparse

from BeautifulSoup import BeautifulSoup
import magicsig
import salmoning
import tornado.escape
import tornado.web

from base import BaseHandler
from logic import content_remote
from logic import smtp
from logic import socialize
from logic import spam
from logic import url_factory
from logic import users

class SalmonHandler(BaseHandler):
  class MockKeyRetriever(magicsig.KeyRetriever):
    def __init__(self, handler, local_user):
      self.handler = handler
      self.local_user = local_user
      self.signer_uri = None

    def LookupPublicKey(self, signer_uri=None):
      self.signer_uri = signer_uri
      user_remote = self.handler.models.users_remote.get(local_username=self.local_user.username, profile_url=signer_uri)[0]

      if not user_remote:
        # get host-meta first
        lrdd_link = users.get_lrdd_link(signer_uri)

        if not lrdd_link:
          raise tornado.web.HTTPError(400)

        # get webfinger
        webfinger_doc = users.get_webfinger(lrdd_link, signer_uri)
        magic_key = webfinger_doc.find('link', rel='magic-public-key')['href'].replace('data:application/magic-public-key,', '')

        if not magic_key:
          raise tornado.web.HTTPError(400)

        new_user = self.handler.models.users_remote()
        new_user.local_username = self.local_user.username
        new_user.profile_url = signer_uri
        new_user.magic_key = magic_key
        new_user.save()

        return (str(magic_key))
      else:
        return (str(user_remote.magic_key))

  def post(self):
    acct = self.get_argument("q").split(':')

    # remove acct: if it's there
    if len(acct) > 1:
      acct[0] = acct[1].split('@')[0]

    acct = acct[0].split('@')[0]

    self.display["acct"] = acct
    user = self.models.users.get(username=acct)[0]

    if not user:
      raise tornado.web.HTTPError(404)

    self.display["user"] = user

    salmonizer = salmoning.SalmonProtocol()
    salmonizer.key_retriever = self.MockKeyRetriever(handler=self, local_user=user)
    env = salmonizer.ParseSalmon(self.request.body, 'application/magic-envelope+xml')
    self.handle_activity_verb(env, salmonizer.key_retriever.signer_uri)

  # tornado & python rock
  def check_xsrf_cookie(self):
    pass

  def handle_activity_verb(self, env, signer_uri):
    salmon_doc = BeautifulSoup(env)
    activity_verb = salmon_doc.find(re.compile('.+:verb$')).string
    user_remote = self.models.users_remote.get(local_username=self.display["user"].username, profile_url=signer_uri)[0]

    if (activity_verb == 'http://activitystrea.ms/schema/1.0/follow'):
      user = users.get_remote_user_info(self, signer_uri, user_remote.local_username)
      user.follower = 1
      user.save()
      smtp.follow(self, user.username, self.display["user"].oauth, user.profile_url)
    elif (activity_verb == 'http://ostatus.org/schema/1.0/unfollow' or activity_verb == 'http://activitystrea.ms/schema/1.0/stop-following'):
      user_remote.follower = 0
      user_remote.save()
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/favorite'):
      activity_object = salmon_doc.find(re.compile('.+:object$'))
      atom_id = activity_object.find('atom:id').string
      local_url = atom_id.split(':')[2]
      content_url = url_factory.load_basic_parameters(self, url=local_url)

      already_favorited = self.models.content_remote.get(to_username=self.display["user"].username,
                                                         from_user=signer_uri,
                                                         type='favorite',
                                                         local_content_name=content_url['name'])[0]
      if already_favorited:
        return

      content = self.models.content.get(username=self.display["user"].username,
                                        name=content_url['name'])[0]
      if not content:
        raise tornado.web.HTTPError(400)
      content.favorites += 1
      content.save()

      favorite_record = self.models.content_remote()
      favorite_record.to_username = self.display["user"].username
      favorite_record.from_user = signer_uri
      favorite_record.username = user_remote.username
      favorite_record.date_created = datetime.datetime.strptime(salmon_doc.find('atom:updated').string[:-6], '%Y-%m-%dT%H:%M:%S')
      favorite_record.type = 'favorite'
      favorite_record.local_content_name = content.name
      favorite_record.save()
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/unfavorite' or activity_verb == 'http://ostatus.org/schema/1.0/unfavorite'):
      # TODO no activity object at least with ostatus??
      pass
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/share'):
      # TODO
      pass
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/post'):
      atom_content = salmon_doc.find('atom:content').string
      sanitized_atom_content = content_remote.sanitize(tornado.escape.xhtml_unescape(atom_content))

      existing_content = self.models.content_remote.get(to_username=self.display["user"].username,
                                                      from_user=signer_uri,
                                                      view=sanitized_atom_content)[0]

      thread = salmon_doc.find('thr:in-reply-to')
      ref = ''
      if thread:
        try:
          local_url = thread['ref'].split(':')[2]
          content_url = url_factory.load_basic_parameters(self, url=local_url)
          content = self.models.content.get(username=self.display["user"].username,
                                            name=content_url['name'])[0]
          if not content:
            raise tornado.web.HTTPError(400)
          ref = content_url['name']
          content.comments_count += 1
          content.comments_updated = datetime.datetime.utcnow()
          content.save()
        except Exception as ex:
          import logging
          logging.error("something wrong with thread")
          logging.error(ex)

      replies = salmon_doc.find('thr:replies')
      comments_count = 0
      comments_updated = None
      if replies:
        if replies.has_key('count'):
          comments_count = replies['count']
        if replies.has_key('updated'):
          comments_updated = replies['updated']
        comments_response = urllib2.urlopen(replies['href'])
        content_remote.parse_feed(self.models, user_remote, comments_response.read(), remote_comments=True)

      mentioned = salmon_doc.findAll('atom:link', rel='mentioned')
      if not mentioned:
        mentioned = salmon_doc.findAll('atom:link', rel='ostatus:attention')

      this_user_mentioned = False
      if mentioned:
        this_user_url = self.nav_url(host=True, username=self.display["user"].username)
        for mentions in mentioned:
          if mentions['href'] == this_user_url:
            # hey, you've been mentioned. cool.
            this_user_mentioned = True
            break

      is_spam = spam.guess(atom_content, self.application.settings["private_path"], self.display["user"].username)

      if existing_content:
        # possible that it's picked up via feed, before we get the salmon call
        post_remote = existing_content
      else:
        post_remote = self.models.content_remote()
      post_remote.to_username = self.display["user"].username
      post_remote.from_user = signer_uri
      post_remote.username = user_remote.username
      post_remote.avatar = user_remote.avatar
      post_remote.date_created = datetime.datetime.strptime(salmon_doc.find('atom:published').string[:-6], '%Y-%m-%dT%H:%M:%S')
      post_remote.date_updated = datetime.datetime.strptime(salmon_doc.find('atom:updated').string[:-6], '%Y-%m-%dT%H:%M:%S')
      post_remote.comments_count = comments_count
      if comments_updated:
        post_remote.comments_updated = datetime.datetime.strptime(comments_updated[:-6], '%Y-%m-%dT%H:%M:%S')
      else:
        post_remote.comments_updated = None
      if is_spam:
        post_remote.is_spam = True
      else:
        spam.train_ham(atom_content, self.application.settings["private_path"], self.display["user"].username)
      post_remote.type = 'comment' if ref else 'post'
      post_remote.title = salmon_doc.find('atom:title').string
      post_remote.post_id = salmon_doc.find('atom:id').string
      post_remote.link = salmon_doc.find('atom:link', rel='alternate')['href']
      post_remote.local_content_name = content.name if ref else ''
      post_remote.view = sanitized_atom_content
      post_remote.save()

      if ref:
        socialize.socialize(self, content)

      if ref:
        smtp.comment(self, post_remote.username, self.display["user"].oauth, self.content_url(content, host=True), sanitized_atom_content)
      elif this_user_mentioned:
        smtp.comment(self, post_remote.username, self.display["user"].oauth, post_remote.link, sanitized_atom_content, this_user_mentioned=True)
<commit_msg>convert count to int<commit_after>import datetime
import re
import urllib2
import urlparse

from BeautifulSoup import BeautifulSoup
import magicsig
import salmoning
import tornado.escape
import tornado.web

from base import BaseHandler
from logic import content_remote
from logic import smtp
from logic import socialize
from logic import spam
from logic import url_factory
from logic import users

class SalmonHandler(BaseHandler):
  class MockKeyRetriever(magicsig.KeyRetriever):
    def __init__(self, handler, local_user):
      self.handler = handler
      self.local_user = local_user
      self.signer_uri = None

    def LookupPublicKey(self, signer_uri=None):
      self.signer_uri = signer_uri
      user_remote = self.handler.models.users_remote.get(local_username=self.local_user.username, profile_url=signer_uri)[0]

      if not user_remote:
        # get host-meta first
        lrdd_link = users.get_lrdd_link(signer_uri)

        if not lrdd_link:
          raise tornado.web.HTTPError(400)

        # get webfinger
        webfinger_doc = users.get_webfinger(lrdd_link, signer_uri)
        magic_key = webfinger_doc.find('link', rel='magic-public-key')['href'].replace('data:application/magic-public-key,', '')

        if not magic_key:
          raise tornado.web.HTTPError(400)

        new_user = self.handler.models.users_remote()
        new_user.local_username = self.local_user.username
        new_user.profile_url = signer_uri
        new_user.magic_key = magic_key
        new_user.save()

        return (str(magic_key))
      else:
        return (str(user_remote.magic_key))

  def post(self):
    acct = self.get_argument("q").split(':')

    # remove acct: if it's there
    if len(acct) > 1:
      acct[0] = acct[1].split('@')[0]

    acct = acct[0].split('@')[0]

    self.display["acct"] = acct
    user = self.models.users.get(username=acct)[0]

    if not user:
      raise tornado.web.HTTPError(404)

    self.display["user"] = user

    salmonizer = salmoning.SalmonProtocol()
    salmonizer.key_retriever = self.MockKeyRetriever(handler=self, local_user=user)
    env = salmonizer.ParseSalmon(self.request.body, 'application/magic-envelope+xml')
    self.handle_activity_verb(env, salmonizer.key_retriever.signer_uri)

  # tornado & python rock
  def check_xsrf_cookie(self):
    pass

  def handle_activity_verb(self, env, signer_uri):
    salmon_doc = BeautifulSoup(env)
    activity_verb = salmon_doc.find(re.compile('.+:verb$')).string
    user_remote = self.models.users_remote.get(local_username=self.display["user"].username, profile_url=signer_uri)[0]

    if (activity_verb == 'http://activitystrea.ms/schema/1.0/follow'):
      user = users.get_remote_user_info(self, signer_uri, user_remote.local_username)
      user.follower = 1
      user.save()
      smtp.follow(self, user.username, self.display["user"].oauth, user.profile_url)
    elif (activity_verb == 'http://ostatus.org/schema/1.0/unfollow' or activity_verb == 'http://activitystrea.ms/schema/1.0/stop-following'):
      user_remote.follower = 0
      user_remote.save()
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/favorite'):
      activity_object = salmon_doc.find(re.compile('.+:object$'))
      atom_id = activity_object.find('atom:id').string
      local_url = atom_id.split(':')[2]
      content_url = url_factory.load_basic_parameters(self, url=local_url)

      already_favorited = self.models.content_remote.get(to_username=self.display["user"].username,
                                                         from_user=signer_uri,
                                                         type='favorite',
                                                         local_content_name=content_url['name'])[0]
      if already_favorited:
        return

      content = self.models.content.get(username=self.display["user"].username,
                                        name=content_url['name'])[0]
      if not content:
        raise tornado.web.HTTPError(400)
      content.favorites += 1
      content.save()

      favorite_record = self.models.content_remote()
      favorite_record.to_username = self.display["user"].username
      favorite_record.from_user = signer_uri
      favorite_record.username = user_remote.username
      favorite_record.date_created = datetime.datetime.strptime(salmon_doc.find('atom:updated').string[:-6], '%Y-%m-%dT%H:%M:%S')
      favorite_record.type = 'favorite'
      favorite_record.local_content_name = content.name
      favorite_record.save()
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/unfavorite' or activity_verb == 'http://ostatus.org/schema/1.0/unfavorite'):
      # TODO no activity object at least with ostatus??
      pass
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/share'):
      # TODO
      pass
    elif (activity_verb == 'http://activitystrea.ms/schema/1.0/post'):
      atom_content = salmon_doc.find('atom:content').string
      sanitized_atom_content = content_remote.sanitize(tornado.escape.xhtml_unescape(atom_content))

      existing_content = self.models.content_remote.get(to_username=self.display["user"].username,
                                                      from_user=signer_uri,
                                                      view=sanitized_atom_content)[0]

      thread = salmon_doc.find('thr:in-reply-to')
      ref = ''
      if thread:
        try:
          local_url = thread['ref'].split(':')[2]
          content_url = url_factory.load_basic_parameters(self, url=local_url)
          content = self.models.content.get(username=self.display["user"].username,
                                            name=content_url['name'])[0]
          if not content:
            raise tornado.web.HTTPError(400)
          ref = content_url['name']
          content.comments_count += 1
          content.comments_updated = datetime.datetime.utcnow()
          content.save()
        except Exception as ex:
          import logging
          logging.error("something wrong with thread")
          logging.error(ex)

      replies = salmon_doc.find('thr:replies')
      comments_count = 0
      comments_updated = None
      if replies:
        if replies.has_key('count'):
          comments_count = int(replies['count'])
        if replies.has_key('updated'):
          comments_updated = replies['updated']
        comments_response = urllib2.urlopen(replies['href'])
        content_remote.parse_feed(self.models, user_remote, comments_response.read(), remote_comments=True)

      mentioned = salmon_doc.findAll('atom:link', rel='mentioned')
      if not mentioned:
        mentioned = salmon_doc.findAll('atom:link', rel='ostatus:attention')

      this_user_mentioned = False
      if mentioned:
        this_user_url = self.nav_url(host=True, username=self.display["user"].username)
        for mentions in mentioned:
          if mentions['href'] == this_user_url:
            # hey, you've been mentioned. cool.
            this_user_mentioned = True
            break

      is_spam = spam.guess(atom_content, self.application.settings["private_path"], self.display["user"].username)

      if existing_content:
        # possible that it's picked up via feed, before we get the salmon call
        post_remote = existing_content
      else:
        post_remote = self.models.content_remote()
      post_remote.to_username = self.display["user"].username
      post_remote.from_user = signer_uri
      post_remote.username = user_remote.username
      post_remote.avatar = user_remote.avatar
      post_remote.date_created = datetime.datetime.strptime(salmon_doc.find('atom:published').string[:-6], '%Y-%m-%dT%H:%M:%S')
      post_remote.date_updated = datetime.datetime.strptime(salmon_doc.find('atom:updated').string[:-6], '%Y-%m-%dT%H:%M:%S')
      post_remote.comments_count = comments_count
      if comments_updated:
        post_remote.comments_updated = datetime.datetime.strptime(comments_updated[:-6], '%Y-%m-%dT%H:%M:%S')
      else:
        post_remote.comments_updated = None
      if is_spam:
        post_remote.is_spam = True
      else:
        spam.train_ham(atom_content, self.application.settings["private_path"], self.display["user"].username)
      post_remote.type = 'comment' if ref else 'post'
      post_remote.title = salmon_doc.find('atom:title').string
      post_remote.post_id = salmon_doc.find('atom:id').string
      post_remote.link = salmon_doc.find('atom:link', rel='alternate')['href']
      post_remote.local_content_name = content.name if ref else ''
      post_remote.view = sanitized_atom_content
      post_remote.save()

      if ref:
        socialize.socialize(self, content)

      if ref:
        smtp.comment(self, post_remote.username, self.display["user"].oauth, self.content_url(content, host=True), sanitized_atom_content)
      elif this_user_mentioned:
        smtp.comment(self, post_remote.username, self.display["user"].oauth, post_remote.link, sanitized_atom_content, this_user_mentioned=True)
<|endoftext|>


<commit_before>"""
Types used to represent a full function/module as an Abstract Syntax Tree.

Most types are small, and are merely used as tokens in the AST. A tree diagram
has been included below to illustrate the relationships between the AST types.

AST Type Tree
-------------

*Basic*
     |--->Assignment
     |             |--->AugmentedAssignment
     |                                    |--->AddAugmentedAssignment
     |                                    |--->SubAugmentedAssignment
     |                                    |--->MulAugmentedAssignment
     |                                    |--->DivAugmentedAssignment
     |                                    |--->ModAugmentedAssignment
     |--->For
"""

from __future__ import print_function, division

from sympy.core import Symbol, Tuple
from sympy.core.basic import Basic
from sympy.core.sympify import _sympify
from sympy.core.relational import Relational
from sympy.utilities.iterables import iterable

class Assignment(Relational):
    """
    Represents variable assignment for code generation.

    Parameters
    ----------
    lhs : Expr
        Sympy object representing the lhs of the expression. These should be
        singular objects, such as one would use in writing code. Notable types
        include Symbol, MatrixSymbol, MatrixElement, and Indexed. Types that
        subclass these types are also supported.

    rhs : Expr
        Sympy object representing the rhs of the expression. This can be any
        type, provided its shape corresponds to that of the lhs. For example,
        a Matrix type can be assigned to MatrixSymbol, but not to Symbol, as
        the dimensions will not align.

    Examples
    ========

    >>> from sympy import symbols, MatrixSymbol, Matrix
    >>> from sympy.codegen.ast import Assignment
    >>> x, y, z = symbols('x, y, z')
    >>> Assignment(x, y)
    Assignment(x, y)
    >>> Assignment(x, 0)
    Assignment(x, 0)
    >>> A = MatrixSymbol('A', 1, 3)
    >>> mat = Matrix([x, y, z]).T
    >>> Assignment(A, mat)
    Assignment(A, Matrix([[x, y, z]]))
    >>> Assignment(A[0, 1], x)
    Assignment(A[0, 1], x)
    """

    rel_op = ':='
    __slots__ = []

    def __new__(cls, lhs, rhs=0, **assumptions):
        from sympy.matrices.expressions.matexpr import (
            MatrixElement, MatrixSymbol)
        from sympy.tensor.indexed import Indexed
        lhs = _sympify(lhs)
        rhs = _sympify(rhs)
        # Tuple of things that can be on the lhs of an assignment
        assignable = (Symbol, MatrixSymbol, MatrixElement, Indexed)
        if not isinstance(lhs, assignable):
            raise TypeError("Cannot assign to lhs of type %s." % type(lhs))
        # Indexed types implement shape, but don't define it until later. This
        # causes issues in assignment validation. For now, matrices are defined
        # as anything with a shape that is not an Indexed
        lhs_is_mat = hasattr(lhs, 'shape') and not isinstance(lhs, Indexed)
        rhs_is_mat = hasattr(rhs, 'shape') and not isinstance(rhs, Indexed)
        # If lhs and rhs have same structure, then this assignment is ok
        if lhs_is_mat:
            if not rhs_is_mat:
                raise ValueError("Cannot assign a scalar to a matrix.")
            elif lhs.shape != rhs.shape:
                raise ValueError("Dimensions of lhs and rhs don't align.")
        elif rhs_is_mat and not lhs_is_mat:
            raise ValueError("Cannot assign a matrix to a scalar.")
        return Relational.__new__(cls, lhs, rhs, **assumptions)

# XXX: This should be handled better
Relational.ValidRelationOperator[':='] = Assignment

class AugmentedAssignment(Assignment):
    """
    Base class for augmented assignments
    """

    @property
    def rel_op(self):
        return self._symbol + '='

class AddAugmentedAssignment(AugmentedAssignment):
    _symbol = '+'

class SubAugmentedAssignment(AugmentedAssignment):
    _symbol = '-'

class MulAugmentedAssignment(AugmentedAssignment):
    _symbol = '*'

class DivAugmentedAssignment(AugmentedAssignment):
    _symbol = '/'

class ModAugmentedAssignment(AugmentedAssignment):
    _symbol = '%'

Relational.ValidRelationOperator['+='] = AddAugmentedAssignment
Relational.ValidRelationOperator['-='] = SubAugmentedAssignment
Relational.ValidRelationOperator['*='] = MulAugmentedAssignment
Relational.ValidRelationOperator['/='] = DivAugmentedAssignment
Relational.ValidRelationOperator['%='] = ModAugmentedAssignment

def aug_assign(lhs, op, rhs):
    """
    Create 'lhs op= rhs'.

    Represents augmented variable assignment for code generation. This is a
    convenience function. You can also use the AugmentedAssignment classes
    directly, like AddAugmentedAssignment(x, y).

    Parameters
    ----------
    lhs : Expr
        Sympy object representing the lhs of the expression. These should be
        singular objects, such as one would use in writing code. Notable types
        include Symbol, MatrixSymbol, MatrixElement, and Indexed. Types that
        subclass these types are also supported.

    op : str
        Operator (+, -, /, *, %).

    rhs : Expr
        Sympy object representing the rhs of the expression. This can be any
        type, provided its shape corresponds to that of the lhs. For example,
        a Matrix type can be assigned to MatrixSymbol, but not to Symbol, as
        the dimensions will not align.

    Examples
    --------

    >>> from sympy import symbols
    >>> from sympy.codegen.ast import aug_assign
    >>> x, y = symbols('x, y')
    >>> aug_assign(x, '+', y)
    AddAugmentedAssignment(x, y)
    """
    if op + '=' not in Relational.ValidRelationOperator:
        raise ValueError("Unrecognized operator %s" % op)
    return Relational.ValidRelationOperator[op + '='](lhs, rhs)

class CodeBlock(Basic):
    """
    Represents a block of code

    For now only assignments are supported. This restriction will be lifted in
    the future.

    Useful methods on this object are

    ``left_hand_sides``: Tuple of left-hand sides of assignments, in order.
    ``left_hand_sides``: Tuple of right-hand sides of assignments, in order.
    ``topological_sort``: Class method. Return a CodeBlock with assignments
                          sorted so that variables are assigned before they
                          are used.
    ``cse``: Return a new CodeBlock with common subexpressions eliminated and
             pulled out as assignments.

    Example
    =======

    >>> from sympy import symbols, ccode
    >>> from sympy.codegen.ast import CodeBlock, Assignment
    >>> x, y = symbols('x y')

    >>> c = CodeBlock(Assignment(x, 1), Assignment(y, x + 1))
    >>> print(ccode(c))
    x = 1;
    y = x + 1;
    """
    def __new__(cls, *args):
        left_hand_sides = []
        right_hand_sides = []
        for i in args:
            if isinstance(i, Assignment):
                lhs, rhs = i.args
                left_hand_sides.append(lhs)
                right_hand_sides.append(rhs)

        obj = Basic.__new__(cls, *args)

        obj.left_hand_sides = Tuple(*left_hand_sides)
        obj.right_hand_sides = Tuple(*right_hand_sides)

        return obj

    @classmethod
    def topological_sort(cls, assignments):
        """
        Return a CodeBlock with topologically sorted assignments so that
        variables are assigned before they are used.

        The existing order of assignments is preserved as much as possible.

        This function assumes that variables are assigned to only once.

        This is a class constructor so that the default constructor for
        CodeBlock can error when variables are used before they are assigned.

        Example
        =======

        >>> from sympy import symbols
        >>> from sympy.codegen.ast import CodeBlock, Assignment
        >>> x, y, z = symbols('x y z')

        >>> assignments = [
        ...     Assignment(x, y + z),
        ...     Assignment(y, z + 1),
        ...     Assignment(z, 2),
        ... ]
        >>> CodeBlock.topological_sort(assignments)
        CodeBlock(Assignment(z, 2), Assignment(y, z + 1), Assignment(x, y + z))
        """
        from sympy.utilities.iterables import topological_sort
        # Create a graph where the nodes are assignments and there is a directed edge
        # between nodes that use a variable and nodes that assign that
        # variable, like

        # [(x := 1, y := x + 1), (x := 1, z := y + z), (y := x + 1, z := y + z)]

        # If we then topologically sort these nodes, they will be in
        # assignment order, like

        # x := 1
        # y := x + 1
        # z := y + z

        # A = The nodes
        #
        # enumerate keeps nodes in the same order they are already in if
        # possible. It will also allow us to handle duplicate assignments to
        # the same variable when those are implemented.
        A = list(enumerate(assignments))

        # var_map = {variable: [assignments using variable]}
        # like {x: [y := x + 1, z := y + x], ...}
        var_map = {}

        # E = Edges in the graph
        E = []
        for i in A:
            if i[1].lhs in var_map:
                E.append((var_map[i[1].lhs], i))
            var_map[i[1].lhs] = i
        for i in A:
            for x in i[1].rhs.free_symbols:
                if x not in var_map:
                    # XXX: Allow this case?
                    raise ValueError("Undefined variable %s" % x)
                E.append((var_map[x], i))

        ordered_assignments = topological_sort([A, E])
        # De-enumerate the result
        return cls(*list(zip(*ordered_assignments))[1])

    def cse(self, symbols=None, optimizations=None, postprocess=None,
        order='canonical'):
        """
        Return a new code block with common subexpressions eliminated

        See the docstring of :func:`sympy.simplify.cse_main.cse` for more
        information.

        Examples
        ========

        >>> from sympy import symbols, sin
        >>> from sympy.codegen.ast import CodeBlock, Assignment
        >>> x, y, z = symbols('x y z')

        >>> c = CodeBlock(
        ...     Assignment(x, 1),
        ...     Assignment(y, sin(x) + 1),
        ...     Assignment(z, sin(x) - 1),
        ... )
        ...
        >>> c.cse()
        CodeBlock(Assignment(x, 1), Assignment(x0, sin(x)), Assignment(y, x0 +
        1), Assignment(z, x0 - 1))
        """
        # TODO: Check that the symbols are new
        from sympy.simplify.cse_main import cse

        if not all(isinstance(i, Assignment) for i in self.args):
            # Will support more things later
            raise NotImplementedError("CodeBlock.cse only supports Assignments")

        if any(isinstance(i, AugmentedAssignment) for i in self.args):
            raise NotImplementedError("CodeBlock.cse doesn't yet work with AugmentedAssignments")

        for i, lhs in enumerate(self.left_hand_sides):
            if lhs in self.left_hand_sides[:i]:
                raise NotImplementedError("Duplicate assignments to the same "
                    "variable are not yet supported (%s)" % lhs)

        replacements, reduced_exprs = cse(self.right_hand_sides, symbols=symbols,
            optimizations=optimizations, postprocess=postprocess, order=order)
        assert len(reduced_exprs) == 1
        new_block = tuple(Assignment(var, expr) for var, expr in
            zip(self.left_hand_sides, reduced_exprs[0]))
        new_assignments = tuple(Assignment(*i) for i in replacements)
        return self.topological_sort(new_assignments + new_block)

class For(Basic):
    """Represents a 'for-loop' in the code.

    Expressions are of the form:
        "for target in iter:
            body..."

    Parameters
    ----------
    target : symbol
    iter : iterable
    body : sympy expr

    Examples
    --------

    >>> from sympy import symbols, Range
    >>> from sympy.codegen.ast import aug_assign, For
    >>> x, n = symbols('x n')
    >>> For(n, Range(10), [aug_assign(x, '+', n)])
    For(n, Range(0, 10, 1), CodeBlock(AddAugmentedAssignment(x, n)))
    """

    def __new__(cls, target, iter, body):
        target = _sympify(target)
        if not iterable(iter):
            raise TypeError("iter must be an iterable")
        if isinstance(iter, list):
            # _sympify errors on lists because they are mutable
            iter = tuple(iter)
        iter = _sympify(iter)
        if not isinstance(body, CodeBlock):
            if not iterable(body):
                raise TypeError("body must be an iterable or CodeBlock")
            body = CodeBlock(*(_sympify(i) for i in body))
        return Basic.__new__(cls, target, iter, body)

    @property
    def target(self):
        return self._args[0]

    @property
    def iterable(self):
        return self._args[1]

    @property
    def body(self):
        return self._args[2]
<commit_msg>Add CodeBlock to the docstring of sympy.codegen.ast<commit_after>"""
Types used to represent a full function/module as an Abstract Syntax Tree.

Most types are small, and are merely used as tokens in the AST. A tree diagram
has been included below to illustrate the relationships between the AST types.

AST Type Tree
-------------

*Basic*
     |--->Assignment
     |             |--->AugmentedAssignment
     |                                    |--->AddAugmentedAssignment
     |                                    |--->SubAugmentedAssignment
     |                                    |--->MulAugmentedAssignment
     |                                    |--->DivAugmentedAssignment
     |                                    |--->ModAugmentedAssignment
     |
     |--->CodeBlock
     |
     |--->For
"""

from __future__ import print_function, division

from sympy.core import Symbol, Tuple
from sympy.core.basic import Basic
from sympy.core.sympify import _sympify
from sympy.core.relational import Relational
from sympy.utilities.iterables import iterable

class Assignment(Relational):
    """
    Represents variable assignment for code generation.

    Parameters
    ----------
    lhs : Expr
        Sympy object representing the lhs of the expression. These should be
        singular objects, such as one would use in writing code. Notable types
        include Symbol, MatrixSymbol, MatrixElement, and Indexed. Types that
        subclass these types are also supported.

    rhs : Expr
        Sympy object representing the rhs of the expression. This can be any
        type, provided its shape corresponds to that of the lhs. For example,
        a Matrix type can be assigned to MatrixSymbol, but not to Symbol, as
        the dimensions will not align.

    Examples
    ========

    >>> from sympy import symbols, MatrixSymbol, Matrix
    >>> from sympy.codegen.ast import Assignment
    >>> x, y, z = symbols('x, y, z')
    >>> Assignment(x, y)
    Assignment(x, y)
    >>> Assignment(x, 0)
    Assignment(x, 0)
    >>> A = MatrixSymbol('A', 1, 3)
    >>> mat = Matrix([x, y, z]).T
    >>> Assignment(A, mat)
    Assignment(A, Matrix([[x, y, z]]))
    >>> Assignment(A[0, 1], x)
    Assignment(A[0, 1], x)
    """

    rel_op = ':='
    __slots__ = []

    def __new__(cls, lhs, rhs=0, **assumptions):
        from sympy.matrices.expressions.matexpr import (
            MatrixElement, MatrixSymbol)
        from sympy.tensor.indexed import Indexed
        lhs = _sympify(lhs)
        rhs = _sympify(rhs)
        # Tuple of things that can be on the lhs of an assignment
        assignable = (Symbol, MatrixSymbol, MatrixElement, Indexed)
        if not isinstance(lhs, assignable):
            raise TypeError("Cannot assign to lhs of type %s." % type(lhs))
        # Indexed types implement shape, but don't define it until later. This
        # causes issues in assignment validation. For now, matrices are defined
        # as anything with a shape that is not an Indexed
        lhs_is_mat = hasattr(lhs, 'shape') and not isinstance(lhs, Indexed)
        rhs_is_mat = hasattr(rhs, 'shape') and not isinstance(rhs, Indexed)
        # If lhs and rhs have same structure, then this assignment is ok
        if lhs_is_mat:
            if not rhs_is_mat:
                raise ValueError("Cannot assign a scalar to a matrix.")
            elif lhs.shape != rhs.shape:
                raise ValueError("Dimensions of lhs and rhs don't align.")
        elif rhs_is_mat and not lhs_is_mat:
            raise ValueError("Cannot assign a matrix to a scalar.")
        return Relational.__new__(cls, lhs, rhs, **assumptions)

# XXX: This should be handled better
Relational.ValidRelationOperator[':='] = Assignment

class AugmentedAssignment(Assignment):
    """
    Base class for augmented assignments
    """

    @property
    def rel_op(self):
        return self._symbol + '='

class AddAugmentedAssignment(AugmentedAssignment):
    _symbol = '+'

class SubAugmentedAssignment(AugmentedAssignment):
    _symbol = '-'

class MulAugmentedAssignment(AugmentedAssignment):
    _symbol = '*'

class DivAugmentedAssignment(AugmentedAssignment):
    _symbol = '/'

class ModAugmentedAssignment(AugmentedAssignment):
    _symbol = '%'

Relational.ValidRelationOperator['+='] = AddAugmentedAssignment
Relational.ValidRelationOperator['-='] = SubAugmentedAssignment
Relational.ValidRelationOperator['*='] = MulAugmentedAssignment
Relational.ValidRelationOperator['/='] = DivAugmentedAssignment
Relational.ValidRelationOperator['%='] = ModAugmentedAssignment

def aug_assign(lhs, op, rhs):
    """
    Create 'lhs op= rhs'.

    Represents augmented variable assignment for code generation. This is a
    convenience function. You can also use the AugmentedAssignment classes
    directly, like AddAugmentedAssignment(x, y).

    Parameters
    ----------
    lhs : Expr
        Sympy object representing the lhs of the expression. These should be
        singular objects, such as one would use in writing code. Notable types
        include Symbol, MatrixSymbol, MatrixElement, and Indexed. Types that
        subclass these types are also supported.

    op : str
        Operator (+, -, /, *, %).

    rhs : Expr
        Sympy object representing the rhs of the expression. This can be any
        type, provided its shape corresponds to that of the lhs. For example,
        a Matrix type can be assigned to MatrixSymbol, but not to Symbol, as
        the dimensions will not align.

    Examples
    --------

    >>> from sympy import symbols
    >>> from sympy.codegen.ast import aug_assign
    >>> x, y = symbols('x, y')
    >>> aug_assign(x, '+', y)
    AddAugmentedAssignment(x, y)
    """
    if op + '=' not in Relational.ValidRelationOperator:
        raise ValueError("Unrecognized operator %s" % op)
    return Relational.ValidRelationOperator[op + '='](lhs, rhs)

class CodeBlock(Basic):
    """
    Represents a block of code

    For now only assignments are supported. This restriction will be lifted in
    the future.

    Useful methods on this object are

    ``left_hand_sides``: Tuple of left-hand sides of assignments, in order.
    ``left_hand_sides``: Tuple of right-hand sides of assignments, in order.
    ``topological_sort``: Class method. Return a CodeBlock with assignments
                          sorted so that variables are assigned before they
                          are used.
    ``cse``: Return a new CodeBlock with common subexpressions eliminated and
             pulled out as assignments.

    Example
    =======

    >>> from sympy import symbols, ccode
    >>> from sympy.codegen.ast import CodeBlock, Assignment
    >>> x, y = symbols('x y')

    >>> c = CodeBlock(Assignment(x, 1), Assignment(y, x + 1))
    >>> print(ccode(c))
    x = 1;
    y = x + 1;
    """
    def __new__(cls, *args):
        left_hand_sides = []
        right_hand_sides = []
        for i in args:
            if isinstance(i, Assignment):
                lhs, rhs = i.args
                left_hand_sides.append(lhs)
                right_hand_sides.append(rhs)

        obj = Basic.__new__(cls, *args)

        obj.left_hand_sides = Tuple(*left_hand_sides)
        obj.right_hand_sides = Tuple(*right_hand_sides)

        return obj

    @classmethod
    def topological_sort(cls, assignments):
        """
        Return a CodeBlock with topologically sorted assignments so that
        variables are assigned before they are used.

        The existing order of assignments is preserved as much as possible.

        This function assumes that variables are assigned to only once.

        This is a class constructor so that the default constructor for
        CodeBlock can error when variables are used before they are assigned.

        Example
        =======

        >>> from sympy import symbols
        >>> from sympy.codegen.ast import CodeBlock, Assignment
        >>> x, y, z = symbols('x y z')

        >>> assignments = [
        ...     Assignment(x, y + z),
        ...     Assignment(y, z + 1),
        ...     Assignment(z, 2),
        ... ]
        >>> CodeBlock.topological_sort(assignments)
        CodeBlock(Assignment(z, 2), Assignment(y, z + 1), Assignment(x, y + z))
        """
        from sympy.utilities.iterables import topological_sort
        # Create a graph where the nodes are assignments and there is a directed edge
        # between nodes that use a variable and nodes that assign that
        # variable, like

        # [(x := 1, y := x + 1), (x := 1, z := y + z), (y := x + 1, z := y + z)]

        # If we then topologically sort these nodes, they will be in
        # assignment order, like

        # x := 1
        # y := x + 1
        # z := y + z

        # A = The nodes
        #
        # enumerate keeps nodes in the same order they are already in if
        # possible. It will also allow us to handle duplicate assignments to
        # the same variable when those are implemented.
        A = list(enumerate(assignments))

        # var_map = {variable: [assignments using variable]}
        # like {x: [y := x + 1, z := y + x], ...}
        var_map = {}

        # E = Edges in the graph
        E = []
        for i in A:
            if i[1].lhs in var_map:
                E.append((var_map[i[1].lhs], i))
            var_map[i[1].lhs] = i
        for i in A:
            for x in i[1].rhs.free_symbols:
                if x not in var_map:
                    # XXX: Allow this case?
                    raise ValueError("Undefined variable %s" % x)
                E.append((var_map[x], i))

        ordered_assignments = topological_sort([A, E])
        # De-enumerate the result
        return cls(*list(zip(*ordered_assignments))[1])

    def cse(self, symbols=None, optimizations=None, postprocess=None,
        order='canonical'):
        """
        Return a new code block with common subexpressions eliminated

        See the docstring of :func:`sympy.simplify.cse_main.cse` for more
        information.

        Examples
        ========

        >>> from sympy import symbols, sin
        >>> from sympy.codegen.ast import CodeBlock, Assignment
        >>> x, y, z = symbols('x y z')

        >>> c = CodeBlock(
        ...     Assignment(x, 1),
        ...     Assignment(y, sin(x) + 1),
        ...     Assignment(z, sin(x) - 1),
        ... )
        ...
        >>> c.cse()
        CodeBlock(Assignment(x, 1), Assignment(x0, sin(x)), Assignment(y, x0 +
        1), Assignment(z, x0 - 1))
        """
        # TODO: Check that the symbols are new
        from sympy.simplify.cse_main import cse

        if not all(isinstance(i, Assignment) for i in self.args):
            # Will support more things later
            raise NotImplementedError("CodeBlock.cse only supports Assignments")

        if any(isinstance(i, AugmentedAssignment) for i in self.args):
            raise NotImplementedError("CodeBlock.cse doesn't yet work with AugmentedAssignments")

        for i, lhs in enumerate(self.left_hand_sides):
            if lhs in self.left_hand_sides[:i]:
                raise NotImplementedError("Duplicate assignments to the same "
                    "variable are not yet supported (%s)" % lhs)

        replacements, reduced_exprs = cse(self.right_hand_sides, symbols=symbols,
            optimizations=optimizations, postprocess=postprocess, order=order)
        assert len(reduced_exprs) == 1
        new_block = tuple(Assignment(var, expr) for var, expr in
            zip(self.left_hand_sides, reduced_exprs[0]))
        new_assignments = tuple(Assignment(*i) for i in replacements)
        return self.topological_sort(new_assignments + new_block)

class For(Basic):
    """Represents a 'for-loop' in the code.

    Expressions are of the form:
        "for target in iter:
            body..."

    Parameters
    ----------
    target : symbol
    iter : iterable
    body : sympy expr

    Examples
    --------

    >>> from sympy import symbols, Range
    >>> from sympy.codegen.ast import aug_assign, For
    >>> x, n = symbols('x n')
    >>> For(n, Range(10), [aug_assign(x, '+', n)])
    For(n, Range(0, 10, 1), CodeBlock(AddAugmentedAssignment(x, n)))
    """

    def __new__(cls, target, iter, body):
        target = _sympify(target)
        if not iterable(iter):
            raise TypeError("iter must be an iterable")
        if isinstance(iter, list):
            # _sympify errors on lists because they are mutable
            iter = tuple(iter)
        iter = _sympify(iter)
        if not isinstance(body, CodeBlock):
            if not iterable(body):
                raise TypeError("body must be an iterable or CodeBlock")
            body = CodeBlock(*(_sympify(i) for i in body))
        return Basic.__new__(cls, target, iter, body)

    @property
    def target(self):
        return self._args[0]

    @property
    def iterable(self):
        return self._args[1]

    @property
    def body(self):
        return self._args[2]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
# Copyright (c) 2020 ContinualAI Research                                      #
# Copyrights licensed under the CC BY 4.0 License.                             #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 1-05-2020                                                              #
# Author(s): ContinualAI                                                       #
# E-mail: contact@continualai.org                                              #
# Website: clair.continualai.org                                               #
################################################################################

""" Rehearsal Strategy Implementation """

# Python 2-3 compatible
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

from training.strategies.strategy import Strategy
from avalanche.evaluation.eval_protocol import EvalProtocol
from avalanche.evaluation.metrics import ACC
from avalanche.training.utils import pad_data, shuffle_in_unison
import torch
import torch.nn.functional as F
import numpy as np
import copy

def distillation_loss(y_pred, y_teacher, temperature):
    """ Distillation loss. """
    scale = y_teacher.shape[-1]  # kl_div is normalized by element instead of observation
    log_p = F.log_softmax(y_pred / temperature, dim=1)
    q = F.softmax(y_teacher / temperature, dim=1)
    res = scale * F.kl_div(log_p, q, reduction='mean')
    return res

class LearningWithoutForgetting(Strategy):
    def __init__(self, model, classes_per_task, alpha=0.5, distillation_loss_T=2, warmup_epochs=0, optimizer=None,
                 criterion=torch.nn.CrossEntropyLoss(), mb_size=256,
                 train_ep=2, device=None, preproc=None,
                 eval_protocol=EvalProtocol(metrics=[ACC()])):
        """
        Learning without Forgetting Strategy.

        paper: https://arxiv.org/abs/1606.09282
        original implementation (Matlab): https://github.com/lizhitwo/LearningWithoutForgetting
        reference implementation (pytorch): https://github.com/arunmallya/packnet/blob/master/src/lwf.py

        Args:
            classes_per_task:
            alpha: distillation loss coefficient. Can be an integer or a list of values (one for each task).
            distillation_loss_T: distillation loss temperature
            warmup_epochs: number of warmup epochs training only the new parameters.
        """
        super(LearningWithoutForgetting, self).__init__(
            model, optimizer, criterion, mb_size, train_ep, multi_head=False,
            device=device, preproc=preproc, eval_protocol=eval_protocol
        )

        # LwF parameters
        self.classes_per_task = classes_per_task
        self.prev_model = None
        self.distillation_loss_T = distillation_loss_T
        self.alpha = alpha
        self.warmup_epochs = warmup_epochs

    def warmup_train(self):
        """ Train only the new parameters for the first epochs. """
        # add only the last layer to the trainable parameters
        opt = torch.optim.SGD(lr=0.01, params=self.model.classifier.parameters())

        train_x, train_y, it_x_ep = self.preproc_batch_data(self.x, self.y, self.t)
        model = self.model.to(self.device)

        train_x = torch.tensor(train_x, dtype=torch.float)
        train_y = torch.tensor(train_y, dtype=torch.long)
        for ep in range(self.warmup_epochs):
            for it in range(it_x_ep):
                start = it * self.mb_size
                end = (it + 1) * self.mb_size

                self.optimizer.zero_grad()
                x_mb = train_x[start:end].to(self.device)
                y_mb = train_y[start:end].to(self.device)
                logits = model(x_mb)
                # loss computed only on the new classes
                loss = self.criterion(logits[:, self.t*self.classes_per_task:(self.t+1)*self.classes_per_task],
                                      y_mb - self.t*self.classes_per_task)
                loss.backward()
                opt.step()

    def compute_loss(self, logits, y_mb):
        dist_loss = 0
        if self.prev_model is not None:
            y_prev = self.prev_model(self.x_mb).detach()
            loss = self.criterion(logits, y_mb)
            dist_loss += distillation_loss(logits, y_prev, self.distillation_loss_T)

            if isinstance(self.alpha, list):
                loss = loss + self.alpha[self.t] * dist_loss
            else:
                loss = loss + self.alpha * dist_loss
        else:
            loss = self.criterion(logits, y_mb)
        return loss

    def before_train(self):
        self.warmup_train()

    def after_train(self):
        self.prev_model = copy.deepcopy(self.model)

<commit_msg>fix bug in import<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
# Copyright (c) 2020 ContinualAI Research                                      #
# Copyrights licensed under the CC BY 4.0 License.                             #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 1-05-2020                                                              #
# Author(s): ContinualAI                                                       #
# E-mail: contact@continualai.org                                              #
# Website: clair.continualai.org                                               #
################################################################################

""" Rehearsal Strategy Implementation """

# Python 2-3 compatible
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

from avalanche.training.strategies.strategy import Strategy
from avalanche.evaluation.eval_protocol import EvalProtocol
from avalanche.evaluation.metrics import ACC
from avalanche.training.utils import pad_data, shuffle_in_unison
import torch
import torch.nn.functional as F
import numpy as np
import copy

def distillation_loss(y_pred, y_teacher, temperature):
    """ Distillation loss. """
    scale = y_teacher.shape[-1]  # kl_div is normalized by element instead of observation
    log_p = F.log_softmax(y_pred / temperature, dim=1)
    q = F.softmax(y_teacher / temperature, dim=1)
    res = scale * F.kl_div(log_p, q, reduction='mean')
    return res

class LearningWithoutForgetting(Strategy):
    def __init__(self, model, classes_per_task, alpha=0.5, distillation_loss_T=2, warmup_epochs=0, optimizer=None,
                 criterion=torch.nn.CrossEntropyLoss(), mb_size=256,
                 train_ep=2, device=None, preproc=None,
                 eval_protocol=EvalProtocol(metrics=[ACC()])):
        """
        Learning without Forgetting Strategy.

        paper: https://arxiv.org/abs/1606.09282
        original implementation (Matlab): https://github.com/lizhitwo/LearningWithoutForgetting
        reference implementation (pytorch): https://github.com/arunmallya/packnet/blob/master/src/lwf.py

        Args:
            classes_per_task:
            alpha: distillation loss coefficient. Can be an integer or a list of values (one for each task).
            distillation_loss_T: distillation loss temperature
            warmup_epochs: number of warmup epochs training only the new parameters.
        """
        super(LearningWithoutForgetting, self).__init__(
            model, optimizer, criterion, mb_size, train_ep, multi_head=False,
            device=device, preproc=preproc, eval_protocol=eval_protocol
        )

        # LwF parameters
        self.classes_per_task = classes_per_task
        self.prev_model = None
        self.distillation_loss_T = distillation_loss_T
        self.alpha = alpha
        self.warmup_epochs = warmup_epochs

    def warmup_train(self):
        """ Train only the new parameters for the first epochs. """
        # add only the last layer to the trainable parameters
        opt = torch.optim.SGD(lr=0.01, params=self.model.classifier.parameters())

        train_x, train_y, it_x_ep = self.preproc_batch_data(self.x, self.y, self.t)
        model = self.model.to(self.device)

        train_x = torch.tensor(train_x, dtype=torch.float)
        train_y = torch.tensor(train_y, dtype=torch.long)
        for ep in range(self.warmup_epochs):
            for it in range(it_x_ep):
                start = it * self.mb_size
                end = (it + 1) * self.mb_size

                self.optimizer.zero_grad()
                x_mb = train_x[start:end].to(self.device)
                y_mb = train_y[start:end].to(self.device)
                logits = model(x_mb)
                # loss computed only on the new classes
                loss = self.criterion(logits[:, self.t*self.classes_per_task:(self.t+1)*self.classes_per_task],
                                      y_mb - self.t*self.classes_per_task)
                loss.backward()
                opt.step()

    def compute_loss(self, logits, y_mb):
        dist_loss = 0
        if self.prev_model is not None:
            y_prev = self.prev_model(self.x_mb).detach()
            loss = self.criterion(logits, y_mb)
            dist_loss += distillation_loss(logits, y_prev, self.distillation_loss_T)

            if isinstance(self.alpha, list):
                loss = loss + self.alpha[self.t] * dist_loss
            else:
                loss = loss + self.alpha * dist_loss
        else:
            loss = self.criterion(logits, y_mb)
        return loss

    def before_train(self):
        self.warmup_train()

    def after_train(self):
        self.prev_model = copy.deepcopy(self.model)

<|endoftext|>