<commit_before>import os
from tempfile import mkdtemp
from zipfile import ZipFile

import shutil

import fiona
from django.contrib.gis.geos import LinearRing
from django.contrib.gis.geos import Polygon
from django.core.management import BaseCommand
from django.db import transaction
from rasterio.warp import transform_geom

from seedsource.models import SeedZone

WASHINGTON_ZONES_DIR = 'WA_NEW_ZONES'
OREGON_ZONES_DIR = 'Oregon_Seed_Zones'
HISTORIC_ZONES_DIR = 'historic_seed_zones'

WASHINGTON_ZONES = {
    'psme': 'PSME.shp',
    'pico': 'PICO.shp',
    'pipo': 'PIPO.shp',
    'thpl': 'THPL.shp',
    'pimo': 'PIMO.shp'
}

OREGON_ZONES = {
    'psme': 'Douglas_Fir.shp',
    'pico': 'Lodgepole_Pine.shp',
    'pipo': 'Ponderosa_Pine.shp',
    'thpl': 'Western_Red_Cedar.shp',
    'pimo': 'Western_White_Pine.shp'
}

SPECIES_NAMES = {
    'psme': 'Douglas-fir',
    'pico': 'Lodgepole pine',
    'pipo': 'Ponderosa pine',
    'thpl': 'Western redcedar',
    'pimo': 'Wester white pine'
}

HISTORIC_ZONES = 'historic_seed_zones.shp'

class Command(BaseCommand):
    help = 'Loads polygon data from seed zone shape files into the database'

    def add_arguments(self, parser):
        parser.add_argument('zones_file', nargs=1, type=str)

    def _add_zones(self, source, name, species, path, zone_field):
        with fiona.open(path, 'r') as shp:
            for feature in shp:
                try:
                    zone_id = feature['properties'][zone_field]
                except KeyError:
                    zone_id = feature['properties'][zone_field.lower()]

                if zone_id is None:
                    continue

                zone_id = int(zone_id)

                if zone_id == 0:
                    continue

                geometry = transform_geom(shp.crs, {'init': 'EPSG:4326'}, feature['geometry'])
                polygon = Polygon(*[LinearRing(x) for x in geometry['coordinates']])
                SeedZone.objects.create(
                    source=source, name=name.format(zone_id=zone_id, species=SPECIES_NAMES.get(species)),
                    species=species, zone_id=zone_id, polygon=polygon
                )

    def handle(self, zones_file, *args, **options):
        temp_dir = mkdtemp()

        try:
            with ZipFile(zones_file[0]) as zf:
                zf.extractall(temp_dir)

            wa_zones = os.path.join(temp_dir, WASHINGTON_ZONES_DIR)
            or_zones = os.path.join(temp_dir, OREGON_ZONES_DIR)
            historic_zones = os.path.join(temp_dir, HISTORIC_ZONES_DIR)

            message = (
                'WARNING: This will replace all your seed zone records and remove associated transfer limits. '
                'Do you want to continue? [y/n]'
            )
            if input(message).lower() not in {'y', 'yes'}:
                return

            with transaction.atomic():
                SeedZone.objects.all().delete()

                print('Loading Washington seed zones...')

                for species, zone_file in WASHINGTON_ZONES.items():
                    self._add_zones(
                        os.path.join(WASHINGTON_ZONES_DIR, zone_file), 'Washington {species} Zone {zone_id}', species,
                        os.path.join(wa_zones, zone_file), 'ZONE_NO'
                    )

                print('Loading Oregon seed zones...')

                for species, zone_file in OREGON_ZONES.items():
                    self._add_zones(
                        os.path.join(OREGON_ZONES_DIR, zone_file), 'Oregon {species} Zone {zone_id}', species,
                        os.path.join(or_zones, zone_file), 'SZ{}'.format(species).upper()
                    )

                print ('Loading historic seed zones...')

                self._add_zones(
                    os.path.join(HISTORIC_ZONES_DIR, HISTORIC_ZONES), 'Washington / Oregon Historic Zone {zone_id}',
                    'generic', os.path.join(historic_zones, HISTORIC_ZONES), 'SUBJ_FSZ'
                )

        finally:
            try:
                shutil.rmtree(temp_dir)
            except OSError:
                pass
<commit_msg>Update seedzone names<commit_after>import os
from tempfile import mkdtemp
from zipfile import ZipFile

import shutil

import fiona
from django.contrib.gis.geos import LinearRing
from django.contrib.gis.geos import Polygon
from django.core.management import BaseCommand
from django.db import transaction
from rasterio.warp import transform_geom

from seedsource.models import SeedZone

WASHINGTON_ZONES_DIR = 'WA_NEW_ZONES'
OREGON_ZONES_DIR = 'Oregon_Seed_Zones'
HISTORIC_ZONES_DIR = 'historic_seed_zones'

WASHINGTON_ZONES = {
    'psme': ('PSME.shp', 'Washington (2002) Douglas-fir Zone {zone_id}'),
    'pico': ('PICO.shp', 'Washington (2002) lodgepole pine Zone {zone_id}'),
    'pipo': ('PIPO.shp', 'Washington (2002) ponderosa pine Zone {zone_id}'),
    'thpl': ('THPL.shp', 'Washington (2002) western redcedar Zone {zone_id}'),
    'pimo': ('PIMO.shp', 'Washington (2002) western white pine Zone {zone_id}')
}

OREGON_ZONES = {
    'psme': ('Douglas_Fir.shp', 'Oregon (1996) Douglas-fir Zone {zone_id}'),
    'pico': ('Lodgepole_Pine.shp', 'Oregon (1996) lodgepole pine Zone {zone_id}'),
    'pipo': ('Ponderosa_Pine.shp', 'Oregon (1996) ponderosa pine Zone {zone_id}'),
    'thpl': ('Western_Red_Cedar.shp', 'Oregon (1996) western redcedar Zone {zone_id}'),
    'pimo': ('Western_White_Pine.shp', 'Oregon (1996) western white pine Zone {zone_id}')
}

SPECIES_NAMES = {
    'psme': 'Douglas-fir',
    'pico': 'Lodgepole pine',
    'pipo': 'Ponderosa pine',
    'thpl': 'Western redcedar',
    'pimo': 'Wester white pine'
}

HISTORIC_ZONES = 'historic_seed_zones.shp'

WA_HISTORIC = {
    11, 12, 30, 201, 202, 211, 212, 221, 222, 231, 222, 231, 232, 240, 401, 402, 403, 411, 412, 421, 422, 430, 440,
    600, 611, 612, 613, 614, 621, 622, 631, 641, 642, 651, 652, 653, 801, 802, 803, 804, 811, 812, 813, 821, 822, 830,
    841, 842, 851
}
OR_HISTORIC = {
    51, 52, 53, 61, 62, 71, 72, 81, 82, 90, 251, 252, 261, 262, 270, 321, 451, 452, 461, 462, 463, 471, 472, 473, 481,
    482, 483, 491, 492, 493, 501, 502, 511, 512, 661, 662, 671, 672, 673, 674, 675, 681, 682, 690, 701, 702, 703, 711,
    712, 713, 721, 722, 731, 751, 853, 862, 863, 871, 872, 881, 882, 883, 891, 892, 901, 902, 911, 912, 921, 922, 930,
    941, 942, 943, 952
}
WA_OR_HISTORIC = {41, 42, 852, 861}

class Command(BaseCommand):
    help = 'Loads polygon data from seed zone shape files into the database'

    def add_arguments(self, parser):
        parser.add_argument('zones_file', nargs=1, type=str)

    def _add_zones(self, source, name, species, path, zone_field):
        with fiona.open(path, 'r') as shp:
            for feature in shp:
                try:
                    zone_id = feature['properties'][zone_field]
                except KeyError:
                    zone_id = feature['properties'][zone_field.lower()]

                if zone_id is None:
                    continue

                zone_id = int(zone_id)

                if zone_id == 0:
                    continue

                if hasattr(name, '__call__'):
                    zone_name = name(zone_id)
                else:
                    zone_name = name.format(zone_id=zone_id, species=SPECIES_NAMES.get(species))

                geometry = transform_geom(shp.crs, {'init': 'EPSG:4326'}, feature['geometry'])
                polygon = Polygon(*[LinearRing(x) for x in geometry['coordinates']])
                SeedZone.objects.create(
                    source=source, name=zone_name, species=species, zone_id=zone_id, polygon=polygon
                )

    def handle(self, zones_file, *args, **options):
        temp_dir = mkdtemp()

        try:
            with ZipFile(zones_file[0]) as zf:
                zf.extractall(temp_dir)

            wa_zones = os.path.join(temp_dir, WASHINGTON_ZONES_DIR)
            or_zones = os.path.join(temp_dir, OREGON_ZONES_DIR)
            historic_zones = os.path.join(temp_dir, HISTORIC_ZONES_DIR)

            message = (
                'WARNING: This will replace all your seed zone records and remove associated transfer limits. '
                'Do you want to continue? [y/n]'
            )
            if input(message).lower() not in {'y', 'yes'}:
                return

            with transaction.atomic():
                SeedZone.objects.all().delete()

                print('Loading Washington seed zones...')

                for species, (zone_file, name) in WASHINGTON_ZONES.items():
                    self._add_zones(
                        os.path.join(WASHINGTON_ZONES_DIR, zone_file), name, species,
                        os.path.join(wa_zones, zone_file), 'ZONE_NO'
                    )

                print('Loading Oregon seed zones...')

                for species, (zone_file, name) in OREGON_ZONES.items():
                    self._add_zones(
                        os.path.join(OREGON_ZONES_DIR, zone_file), name, species, os.path.join(or_zones, zone_file),
                        'SZ{}'.format(species).upper()
                    )

                print('Loading historic seed zones...')

                def get_historic_name(zone_id):
                    if zone_id in WA_HISTORIC:
                        state = 'Washington'
                    elif zone_id in OR_HISTORIC:
                        state = 'Oregon'
                    elif zone_id in WA_OR_HISTORIC:
                        state = 'Oregon & Washington'
                    else:
                        raise ValueError('Could not find state for zone {}'.format(zone_id))

                    return '{} (1966/1973) Zone {:03d}'.format(state, zone_id)

                self._add_zones(
                    os.path.join(HISTORIC_ZONES_DIR, HISTORIC_ZONES), get_historic_name, 'generic',
                    os.path.join(historic_zones, HISTORIC_ZONES), 'SUBJ_FSZ'
                )

        finally:
            try:
                shutil.rmtree(temp_dir)
            except OSError:
                pass
<|endoftext|>


<commit_before># -*- coding: UTF-8
#   cleaning_sched
#   **************
#
# Implementation of the cleaning operations (delete incomplete submission,
# delete expired tips, etc)

import os
import sys

from twisted.internet.defer import inlineCallbacks

from globaleaks.settings import transact, GLSetting

from globaleaks.utils import log, pretty_date_time, is_expired
from globaleaks.jobs.base import GLJob
from globaleaks.models import InternalTip, ReceiverTip, ReceiverFile, InternalFile, Comment
from datetime import datetime

__all__ = ['APSCleaning']

@transact
def get_tiptime_by_marker(store, marker):
    assert marker in InternalTip._marker

    itip_list = store.find(InternalTip, InternalTip.mark == marker)

    tipinfo_list = []
    for itip in itip_list:

        comment_cnt = store.find(Comment, Comment.internaltip_id == itip.id).count()
        files_cnt = store.find(InternalFile, InternalFile.internaltip_id == itip.id).count()

        serialized_tipinfo = {
            'id': itip.id,
            'creation_date': pretty_date_time(itip.creation_date),
            'tip_life_seconds':  itip.context.tip_timetolive,
            'submission_life_seconds':  itip.context.submission_timetolive,
            'files': files_cnt,
            'comments': comment_cnt,
        }
        tipinfo_list.append(serialized_tipinfo)

    return tipinfo_list

def iso2dateobj(str) :
    try:
        ret = datetime.strptime(str, "%Y-%m-%dT%H:%M:%S")
    except ValueError :
        ret = datetime.strptime(str, "%Y-%m-%dT%H:%M:%S.%f")
        ret.replace(microsecond=0)
    return ret

@transact
def itip_cleaning(store, id):
    """
    @param id: aim for an InternalTip, and delete them.
    """
    tit = store.find(InternalTip, InternalTip.id == id).one()

    if not tit: # gtfo
        log.err("Requested invalid InternalTip id in itip_cleaning! %s" % id)
        return

    comments = store.find(Comment, Comment.internaltip_id == id)
    log.debug("[-] Removing %d comments, %d files from an InternalTip and %d rtips" %
        (tit.internalfiles.count(), comments.count(), tit.receivertips.count() ))

    for ifile in tit.internalfiles:
        abspath = os.path.join(GLSetting.submission_path, ifile.file_path)
        ifname = unicode(ifile.name)

        if not os.path.isfile(abspath):
            log.err("Unable to remove %s not existent file, in itip %s has an internalfile %s(%d) missing on FS" %
                (abspath, id, ifname, ifile.size) )
            continue

        try:
            os.unlink(abspath)
        except OSError as excep:
            log.err("Unable to remove %s: %s" % (abspath, excep.strerror) )

        log.debug("Receiver file associated to %s: %d" % (ifname, tit.internalfiles.count()) )

        for rfile in tit.internalfiles:
            try:
                store.remove(rfile)
            except Exception as excep:
                # This happen only if delete on cascade is working #96
                log.debug("Unable to remove ReceiverFile of %s: skipped" % ifname)
                continue
        try:
            store.remove(ifile)
            log.debug("Removed InternalFile %s" % ifname)
        except Exception as excep:
            log.debug("Unable to remove InternalFile %s: skipped" % ifname)
            continue

    for comment in comments:
        store.remove(comment)

    for rtip in tit.receivertips:
        rname = rtip.receiver.name
        try:
            store.remove(rtip)
            log.debug("removed ReceiverTip of %s" % rname)
        except Exception as excep:
            log.debug("Unable to remove ReceiverTip of %s" % rname)
            continue

    # Finally remove a Tip, better if on cascade works :( #96
    store.remove(tit)

@transact
def debug_count_itips_by_marker(store):
    info_list = []
    for marker in InternalTip._marker:
        single_info = {
            marker: store.find(InternalTip, InternalTip.mark == marker).count()
        }
        info_list.append(single_info)

    return info_list

class APSCleaning(GLJob):

    @inlineCallbacks
    def operation(self):
        """
        Goal of this function is to check all the submission not
        finalized, and, if the expiration time sets in the context has
        been reached, then clean the submission_gus along with the fields,
        and, if present, the uploaded folder/files.

        Second goal of this function, is to check all the internaltips
        and their expiration date, if match, remove that, all the folder,
        comment and tip related.
        """
        try:
            submissions = yield get_tiptime_by_marker(InternalTip._marker[0]) # Submission
            log.debug("(Cleaning routines) %d unfinished Submission are check if expired" % len(submissions))
            for submission in submissions:
                if is_expired(iso2dateobj(submission['creation_date']), seconds=submission['submission_life_seconds']):
                    log.info("Deleting an unfinalized Submission (creation date: %s) files %d" %
                             (submission['creation_date'], submission['files']) )
                    yield itip_cleaning(submission['id'])

            tips = yield get_tiptime_by_marker(InternalTip._marker[2]) # First
            log.debug("(Cleaning routines) %d Tips stored are check if expired" % len(tips))
            for tip in tips:
                if is_expired(iso2dateobj(tip['creation_date']), seconds=tip['tip_life_seconds']):
                    log.info("Deleting an expired Tip (creation date: %s) files %d comments %d" %
                             (tip['creation_date'], tip['files'], tip['comments']) )
                    yield itip_cleaning(tip['id'])
        except:
            sys.excepthook(*sys.exc_info())
<commit_msg>fixed issue #101 (useless when DELETE ON CASCADE will work)<commit_after># -*- coding: UTF-8
#   cleaning_sched
#   **************
#
# Implementation of the cleaning operations (delete incomplete submission,
# delete expired tips, etc)

import os
import sys

from twisted.internet.defer import inlineCallbacks

from globaleaks.settings import transact, GLSetting

from globaleaks.utils import log, pretty_date_time, is_expired
from globaleaks.jobs.base import GLJob
from globaleaks.models import InternalTip, ReceiverTip, ReceiverFile, InternalFile, Comment
from datetime import datetime

__all__ = ['APSCleaning']

@transact
def get_tiptime_by_marker(store, marker):
    assert marker in InternalTip._marker

    itip_list = store.find(InternalTip, InternalTip.mark == marker)

    tipinfo_list = []
    for itip in itip_list:

        comment_cnt = store.find(Comment, Comment.internaltip_id == itip.id).count()
        files_cnt = store.find(InternalFile, InternalFile.internaltip_id == itip.id).count()

        if not itip.context:
            log.err("A Tip related to a not existent Context! This would not happen if delete on cascade is working")
            # And the removal is forced putting 1 second of life to the Tip.
            tip_timetolive = 1
            submission_timetolive = 1
        else:
            tip_timetolive = itip.context.tip_timetolive
            submission_timetolive = itip.context.submission_timetolive

        serialized_tipinfo = {
            'id': itip.id,
            'creation_date': pretty_date_time(itip.creation_date),
            'tip_life_seconds':  tip_timetolive,
            'submission_life_seconds':  submission_timetolive,
            'files': files_cnt,
            'comments': comment_cnt,
        }
        tipinfo_list.append(serialized_tipinfo)

    return tipinfo_list

def iso2dateobj(str) :
    try:
        ret = datetime.strptime(str, "%Y-%m-%dT%H:%M:%S")
    except ValueError :
        ret = datetime.strptime(str, "%Y-%m-%dT%H:%M:%S.%f")
        ret.replace(microsecond=0)
    return ret

@transact
def itip_cleaning(store, id):
    """
    @param id: aim for an InternalTip, and delete them.
    """
    tit = store.find(InternalTip, InternalTip.id == id).one()

    if not tit: # gtfo
        log.err("Requested invalid InternalTip id in itip_cleaning! %s" % id)
        return

    comments = store.find(Comment, Comment.internaltip_id == id)
    log.debug("[-] Removing %d comments, %d files from an InternalTip and %d rtips" %
        (tit.internalfiles.count(), comments.count(), tit.receivertips.count() ))

    for ifile in tit.internalfiles:
        abspath = os.path.join(GLSetting.submission_path, ifile.file_path)
        ifname = unicode(ifile.name)

        if not os.path.isfile(abspath):
            log.err("Unable to remove %s not existent file, in itip %s has an internalfile %s(%d) missing on FS" %
                (abspath, id, ifname, ifile.size) )
            continue

        try:
            os.unlink(abspath)
        except OSError as excep:
            log.err("Unable to remove %s: %s" % (abspath, excep.strerror) )

        log.debug("Receiver file associated to %s: %d" % (ifname, tit.internalfiles.count()) )

        for rfile in tit.internalfiles:
            try:
                store.remove(rfile)
            except Exception as excep:
                # This happen only if delete on cascade is working #96
                log.debug("Unable to remove ReceiverFile of %s: skipped" % ifname)
                continue
        try:
            store.remove(ifile)
            log.debug("Removed InternalFile %s" % ifname)
        except Exception as excep:
            log.debug("Unable to remove InternalFile %s: skipped" % ifname)
            continue

    for comment in comments:
        store.remove(comment)

    for rtip in tit.receivertips:
        rname = rtip.receiver.name
        try:
            store.remove(rtip)
            log.debug("removed ReceiverTip of %s" % rname)
        except Exception as excep:
            log.debug("Unable to remove ReceiverTip of %s" % rname)
            continue

    # Finally remove a Tip, better if on cascade works :( #96
    store.remove(tit)

@transact
def debug_count_itips_by_marker(store):
    info_list = []
    for marker in InternalTip._marker:
        single_info = {
            marker: store.find(InternalTip, InternalTip.mark == marker).count()
        }
        info_list.append(single_info)

    return info_list

class APSCleaning(GLJob):

    @inlineCallbacks
    def operation(self):
        """
        Goal of this function is to check all the submission not
        finalized, and, if the expiration time sets in the context has
        been reached, then clean the submission_gus along with the fields,
        and, if present, the uploaded folder/files.

        Second goal of this function, is to check all the internaltips
        and their expiration date, if match, remove that, all the folder,
        comment and tip related.
        """
        try:
            submissions = yield get_tiptime_by_marker(InternalTip._marker[0]) # Submission
            log.debug("(Cleaning routines) %d unfinished Submission are check if expired" % len(submissions))
            for submission in submissions:
                if is_expired(iso2dateobj(submission['creation_date']), seconds=submission['submission_life_seconds']):
                    log.info("Deleting an unfinalized Submission (creation date: %s) files %d" %
                             (submission['creation_date'], submission['files']) )
                    yield itip_cleaning(submission['id'])

            tips = yield get_tiptime_by_marker(InternalTip._marker[2]) # First
            log.debug("(Cleaning routines) %d Tips stored are check if expired" % len(tips))
            for tip in tips:
                if is_expired(iso2dateobj(tip['creation_date']), seconds=tip['tip_life_seconds']):
                    log.info("Deleting an expired Tip (creation date: %s) files %d comments %d" %
                             (tip['creation_date'], tip['files'], tip['comments']) )
                    yield itip_cleaning(tip['id'])
        except:
            sys.excepthook(*sys.exc_info())
<|endoftext|>


<commit_before># -*- coding: UTF-8
#   cleaning_sched
#   **************
#
# Implementation of the cleaning operations (delete incomplete submission,
# delete expired tips, etc)

import os
import sys

from twisted.internet.defer import inlineCallbacks

from globaleaks.settings import transact, transact_ro, GLSetting
from globaleaks.utils.utility import log, pretty_date_time, is_expired, iso2dateobj
from globaleaks.jobs.base import GLJob
from globaleaks.models import InternalTip, ReceiverFile, InternalFile, Comment

__all__ = ['APSCleaning']

@transact_ro
def get_tiptime_by_marker(store, marker):
    assert marker in InternalTip._marker

    itip_list = store.find(InternalTip, InternalTip.mark == marker)

    tipinfo_list = []
    for itip in itip_list:

        comment_cnt = store.find(Comment, Comment.internaltip_id == itip.id).count()
        files_cnt = store.find(InternalFile, InternalFile.internaltip_id == itip.id).count()

        if not itip.context:
            log.err("A Tip related to a not existent Context! This would not happen if delete on cascade is working")
            # And the removal is forced putting 1 second of life to the Tip.
            tip_timetolive = 1
            submission_timetolive = 1
        else:
            tip_timetolive = itip.context.tip_timetolive
            submission_timetolive = itip.context.submission_timetolive

        serialized_tipinfo = {
            'id': itip.id,
            'creation_date': pretty_date_time(itip.creation_date),
            'tip_life_seconds':  tip_timetolive,
            'submission_life_seconds':  submission_timetolive,
            'files': files_cnt,
            'comments': comment_cnt,
        }
        tipinfo_list.append(serialized_tipinfo)

    return tipinfo_list

@transact
def itip_cleaning(store, id):
    """
    @param id: aim for an InternalTip, and delete them.
    """
    tit = store.find(InternalTip, InternalTip.id == id).one()

    if not tit: # gtfo
        log.err("Requested invalid InternalTip id in itip_cleaning! %s" % id)
        return

    comments = store.find(Comment, Comment.internaltip_id == id)
    log.debug("[-] Removing %d comments, %d files from an InternalTip and %d rtips" %
        (tit.internalfiles.count(), comments.count(), tit.receivertips.count() ))

    for ifile in tit.internalfiles:
        abspath = os.path.join(GLSetting.submission_path, ifile.file_path)
        ifname = unicode(ifile.name)

        if not os.path.isfile(abspath):
            log.err("Unable to remove non existent internalfile %s (itip %s, internalfile %s(%d))" %
                (abspath, id, ifname, ifile.size))
        else:
            try:
                os.remove(abspath)
            except OSError as excep:
                log.err("Unable to remove %s: %s" % (abspath, excep.strerror))

        rfiles = store.find(ReceiverFile, ReceiverFile.internalfile_id == ifile.id)
        for rfile in rfiles:
            # The following code must be bypassed if rfile.file_path == ifile.filepath
            if rfile.file_path == ifile.file_path:
                continue

            abspath = os.path.join(GLSetting.submission_path, rfile.file_path)

            if not os.path.isfile(abspath):
                log.err("Unable to remove non existent receiverfile %s (itip %s, internalfile %s(%d))" %
                    (abspath, id, ifname, ifile.size))
                continue

            try:
                os.remove(abspath)
            except OSError as excep:
                log.err("Unable to remove %s: %s" % (abspath, excep.strerror))

    store.remove(tit)

@transact
def debug_count_itips_by_marker(store):
    info_list = []
    for marker in InternalTip._marker:
        single_info = {
            marker: store.find(InternalTip, InternalTip.mark == marker).count()
        }
        info_list.append(single_info)

    return info_list

class APSCleaning(GLJob):
    @staticmethod
    @inlineCallbacks
    def operation():
        """
        Goal of this function is to check all the submission not
        finalized, and, if the expiration time sets in the context has
        been reached, then clean the submission_gus along with the fields,
        and, if present, the uploaded folder/files.

        Second goal of this function, is to check all the InternalTip(s)
        and their expiration date, if match, remove that, all the folder,
        comment and tip related.

        Thir goal of this function is to reset the exception counter that
        acts as limit for mail storm
        """
        try:
            # First Goal
            submissions = yield get_tiptime_by_marker(InternalTip._marker[0]) # Submission
            log.debug("(Cleaning routines) %d unfinished Submission are check if expired" % len(submissions))
            for submission in submissions:
                if is_expired(iso2dateobj(submission['creation_date'])):
                    log.info("Deleting an unfinalized Submission (creation date: %s) files %d" %
                             (submission['creation_date'], submission['files']) )
                    yield itip_cleaning(submission['id'])

            # Second Goal
            tips = yield get_tiptime_by_marker(InternalTip._marker[2]) # First
            log.debug("(Cleaning routines) %d Tips stored are check if expired" % len(tips))
            for tip in tips:
                if is_expired(iso2dateobj(tip['creation_date'])):
                    log.info("Deleting an expired Tip (creation date: %s) files %d comments %d" %
                             (tip['creation_date'], tip['files'], tip['comments']) )
                    yield itip_cleaning(tip['id'])

            # Third Goal: Reset of GLSetting.exceptions
            GLSetting.exceptions = {}

        except Exception as excep:
            log.err("Exception failure in submission/tip cleaning routine (%s)" % excep.message)
            sys.excepthook(*sys.exc_info())
<commit_msg>fixed a log.debug bug<commit_after># -*- coding: UTF-8
#   cleaning_sched
#   **************
#
# Implementation of the cleaning operations (delete incomplete submission,
# delete expired tips, etc)

import os
import sys

from twisted.internet.defer import inlineCallbacks

from globaleaks.settings import transact, transact_ro, GLSetting
from globaleaks.utils.utility import log, pretty_date_time, is_expired, iso2dateobj
from globaleaks.jobs.base import GLJob
from globaleaks.models import InternalTip, ReceiverFile, InternalFile, Comment

__all__ = ['APSCleaning']

@transact_ro
def get_tiptime_by_marker(store, marker):
    assert marker in InternalTip._marker

    itip_list = store.find(InternalTip, InternalTip.mark == marker)

    tipinfo_list = []
    for itip in itip_list:

        comment_cnt = store.find(Comment, Comment.internaltip_id == itip.id).count()
        files_cnt = store.find(InternalFile, InternalFile.internaltip_id == itip.id).count()

        if not itip.context:
            log.err("A Tip related to a not existent Context! This would not happen if delete on cascade is working")
            # And the removal is forced putting 1 second of life to the Tip.
            tip_timetolive = 1
            submission_timetolive = 1
        else:
            tip_timetolive = itip.context.tip_timetolive
            submission_timetolive = itip.context.submission_timetolive

        serialized_tipinfo = {
            'id': itip.id,
            'creation_date': pretty_date_time(itip.creation_date),
            'tip_life_seconds':  tip_timetolive,
            'submission_life_seconds':  submission_timetolive,
            'files': files_cnt,
            'comments': comment_cnt,
        }
        tipinfo_list.append(serialized_tipinfo)

    return tipinfo_list

@transact
def itip_cleaning(store, id):
    """
    @param id: aim for an InternalTip, and delete them.
    """
    tit = store.find(InternalTip, InternalTip.id == id).one()

    if not tit: # gtfo
        log.err("Requested invalid InternalTip id in itip_cleaning! %s" % id)
        return

    comments = store.find(Comment, Comment.internaltip_id == id)
    log.debug("[-] Removing [%d comments] [%d files] [%d rtips] from an InternalTip" %
        (comments.count(), tit.internalfiles.count(), tit.receivertips.count() ))

    for ifile in tit.internalfiles:
        abspath = os.path.join(GLSetting.submission_path, ifile.file_path)
        ifname = unicode(ifile.name)

        if not os.path.isfile(abspath):
            log.err("Unable to remove non existent internalfile %s (itip %s, internalfile %s(%d))" %
                (abspath, id, ifname, ifile.size))
        else:
            try:
                os.remove(abspath)
            except OSError as excep:
                log.err("Unable to remove %s: %s" % (abspath, excep.strerror))

        rfiles = store.find(ReceiverFile, ReceiverFile.internalfile_id == ifile.id)
        for rfile in rfiles:
            # The following code must be bypassed if rfile.file_path == ifile.filepath
            if rfile.file_path == ifile.file_path:
                continue

            abspath = os.path.join(GLSetting.submission_path, rfile.file_path)

            if not os.path.isfile(abspath):
                log.err("Unable to remove non existent receiverfile %s (itip %s, internalfile %s(%d))" %
                    (abspath, id, ifname, ifile.size))
                continue

            try:
                os.remove(abspath)
            except OSError as excep:
                log.err("Unable to remove %s: %s" % (abspath, excep.strerror))

    store.remove(tit)

@transact
def debug_count_itips_by_marker(store):
    info_list = []
    for marker in InternalTip._marker:
        single_info = {
            marker: store.find(InternalTip, InternalTip.mark == marker).count()
        }
        info_list.append(single_info)

    return info_list

class APSCleaning(GLJob):
    @staticmethod
    @inlineCallbacks
    def operation():
        """
        Goal of this function is to check all the submission not
        finalized, and, if the expiration time sets in the context has
        been reached, then clean the submission_gus along with the fields,
        and, if present, the uploaded folder/files.

        Second goal of this function, is to check all the InternalTip(s)
        and their expiration date, if match, remove that, all the folder,
        comment and tip related.

        Thir goal of this function is to reset the exception counter that
        acts as limit for mail storm
        """
        try:
            # First Goal
            submissions = yield get_tiptime_by_marker(InternalTip._marker[0]) # Submission
            log.debug("(Cleaning routines) %d unfinished Submission are check if expired" % len(submissions))
            for submission in submissions:
                if is_expired(iso2dateobj(submission['creation_date'])):
                    log.info("Deleting an unfinalized Submission (creation date: %s) files %d" %
                             (submission['creation_date'], submission['files']) )
                    yield itip_cleaning(submission['id'])

            # Second Goal
            tips = yield get_tiptime_by_marker(InternalTip._marker[2]) # First
            log.debug("(Cleaning routines) %d Tips stored are check if expired" % len(tips))
            for tip in tips:
                if is_expired(iso2dateobj(tip['creation_date'])):
                    log.info("Deleting an expired Tip (creation date: %s) files %d comments %d" %
                             (tip['creation_date'], tip['files'], tip['comments']) )
                    yield itip_cleaning(tip['id'])

            # Third Goal: Reset of GLSetting.exceptions
            GLSetting.exceptions = {}

        except Exception as excep:
            log.err("Exception failure in submission/tip cleaning routine (%s)" % excep.message)
            sys.excepthook(*sys.exc_info())
<|endoftext|>


<commit_before>import os

__VERSION__ = ''
__PLUGIN_VERSION__ = ''

# Config settings
USERNAME = ''
SECRET = ''
API_KEY = ''

DEBUG = False
SOCK_DEBUG = False

EXPERT_MODE = False

ALERT_ON_MSG = True
LOG_TO_CONSOLE = False

BASE_DIR = os.path.expanduser(os.path.join('~', 'floobits'))

# Shared globals
DEFAULT_HOST = 'floobits.com'
DEFAULT_PORT = 3448
SECURE = True
ERROR_COUNT = 0
ERRORS_SENT = 0
# Don't spam us with error reports
MAX_ERROR_REPORTS = 3

PROXY_PORT = 0  # Random port
SHARE_DIR = None
COLAB_DIR = ''
PROJECT_PATH = ''
WORKSPACE_WINDOW = None

PERMS = []
STALKER_MODE = False
SPLIT_MODE = False

AUTO_GENERATED_ACCOUNT = False
PLUGIN_PATH = None

CHAT_VIEW = None
CHAT_VIEW_PATH = None

TICK_TIME = 100
AGENT = None
IGNORE = None

IGNORE_MODIFIED_EVENTS = False
VIEW_TO_HASH = {}

FLOORC_PATH = os.path.expanduser(os.path.join('~', '.floorc'))
<commit_msg>s/stalk/follow/g<commit_after>import os

__VERSION__ = ''
__PLUGIN_VERSION__ = ''

# Config settings
USERNAME = ''
SECRET = ''
API_KEY = ''

DEBUG = False
SOCK_DEBUG = False

EXPERT_MODE = False

ALERT_ON_MSG = True
LOG_TO_CONSOLE = False

BASE_DIR = os.path.expanduser(os.path.join('~', 'floobits'))

# Shared globals
DEFAULT_HOST = 'floobits.com'
DEFAULT_PORT = 3448
SECURE = True
ERROR_COUNT = 0
ERRORS_SENT = 0
# Don't spam us with error reports
MAX_ERROR_REPORTS = 3

PROXY_PORT = 0  # Random port
SHARE_DIR = None
COLAB_DIR = ''
PROJECT_PATH = ''
WORKSPACE_WINDOW = None

PERMS = []
FOLLOW_MODE = False
SPLIT_MODE = False

AUTO_GENERATED_ACCOUNT = False
PLUGIN_PATH = None

CHAT_VIEW = None
CHAT_VIEW_PATH = None

TICK_TIME = 100
AGENT = None
IGNORE = None

IGNORE_MODIFIED_EVENTS = False
VIEW_TO_HASH = {}

FLOORC_PATH = os.path.expanduser(os.path.join('~', '.floorc'))
<|endoftext|>


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

"""
:mod:`birthdayreminder` --- Birthday reminder plugin
====================================================
"""

from __future__ import unicode_literals

__docformat__ = "restructuredtext en"

import atexit
import datetime
from threading import Timer
import random
import operator

from Skype4Py.enums import cmsReceived

from plugin import Plugin
from output import ChatMessage

INTERJECTIONS = [
    "Hurrah",
    "Hooray",
    "Huzzah",
    "Fancy",
    "Awesome",
    "Whee",
    "Whoa",
    "Unbelievable",
    "Wow",
    "Amazing",
    "Yahoo",
    "Yippie",
    "Yay",
    "Yeah",
    "Exciting",
    "Yee-haw",
    "Sweet",
    "Fantastic",
    "OMG",
]

TODAY = "Today is {name} birthday! {interjection}!"

PREVIOUS = "It was {name} birthday {delta_days} {day_or_days} ago."

UPCOMING = "{name} birthday is coming up next in {delta_days} {day_or_days}."

DAY_OR_DAYS = {
    "singular": "day",
    "plural": "days",
}

POSSESSIVE = "{name}'s"

CONJUNCTION = "and"

def humanize_names_list(names_list,
                        possessive=POSSESSIVE,
                        conjunction=CONJUNCTION):
    """
    >>> possessive = "{name}'s"
    >>> conjunction = "and"

    >>> humanize_names_list(["herp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's"

    >>> humanize_names_list(["herp", "derp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's and derp's"

    >>> humanize_names_list(["herp", "derp", "durp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's, derp's and durp's"
    """

    assert isinstance(names_list, (list, tuple))
    names = [possessive.format(name=name) for name in names_list[:]]
    if len(names) is 1:
        return names[0]
    elif len(names) > 1:
        return "{0} {conjunction} {1}".format(", ".join(names[:-1]), names[-1],
                                              conjunction=conjunction)

def str_to_datetime(date_str):
    """
    >>> str_to_datetime("13.1")
    datetime.datetime(1900, 1, 13, 0, 0)

    >>> str_to_datetime("13.01")
    datetime.datetime(1900, 1, 13, 0, 0)

    >>> str_to_datetime("01.13")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format

    >>> str_to_datetime("13.1.2000")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("2000-01-13")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("2000-13-01")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format

    >>> str_to_datetime("13.01.2000")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("01.13.2000")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format
    """

    dt = None
    formats = ["%d.%m", "%d.%m.%Y", "%Y-%m-%d"]
    for fmt in formats:
        try:
            dt = datetime.datetime.strptime(date_str, fmt)
            break
        except ValueError:
            pass
    if dt is None:
        raise ValueError("Input string does not match any known date format")
    return dt

_timer = None

class BirthdayReminder(Plugin):
    # Time in seconds.
    CHECK_INTERVAL = 600

    @staticmethod
    @atexit.register
    def _cleanup():
        global _timer
        try:
            _timer.cancel()
            _timer.join()
        except AttributeError:
            pass
        finally:
            del _timer

    def __init__(self, priority=0, whitelist=None, **kwargs):
        super(BirthdayReminder, self).__init__(priority, whitelist, **kwargs)

        self.dates = dict()
        birthdays = self.options.setdefault("birthdays")
        for name, date_str in birthdays.iteritems():
            try:
                dt = str_to_datetime(date_str)
            except ValueError:
                self._logger.error("Parse error: '%s' for '%s'", name, date_str)
            else:
                self.dates.update({name: dt})

        if not self.dates:
            self._logger.warning("No valid dates has been parsed!")
            return

        self._notified = False
        self._logger.info("Parsed %d date(s)", len(self.dates))
        self._check_dates()

    def _check_dates(self):
        global _timer

        today = datetime.datetime.today()

        persons = list()
        for name, dt in self.dates.iteritems():
            dt = dt.replace(year=today.year)
            delta = today - dt
            if delta.days == 0:
                persons.append(name)
        if persons:
            if not self._notified:
                self._notify_chats(persons)
                self._notified = True
            return
        self._notified = False
        _timer = Timer(self.CHECK_INTERVAL, self._check_dates)
        _timer.daemon = True
        _timer.start()

    def _notify_chats(self, names):
        substitutes = {
            "name": humanize_names_list(names),
            "interjection": random.choice(INTERJECTIONS),
        }
        message = TODAY.format(**substitutes)

        for chat in self.whitelist:
            self.output.append(ChatMessage(chat, message))

    def on_message_status(self, message, status):
        if status != cmsReceived:
            return

        if not message.Body.strip().startswith("!birthdays"):
            return

        today = datetime.datetime.today()

        previous_persons = list()
        upcoming_persons = list()
        today_names = list()
        for name, dt in self.dates.iteritems():
            dt = dt.replace(year=today.year)
            delta = today - dt
            person = (name, delta)
            if delta.days > 0:
                previous_persons.append(person)
            elif delta.days < 0:
                upcoming_persons.append(person)
            elif delta.days == 0:
                today_names.append(name)

        def __same_by_delta(persons, delta):
            for _name, _delta in persons:
                if _delta.days == delta.days:
                    yield (_name, _delta)

        previous_names = list()
        previous_person = min(previous_persons, key=operator.itemgetter(1))
        previous_delta = previous_person[1]
        previous_delta_days = previous_delta.days
        for name, _ in __same_by_delta(previous_persons, previous_delta):
            previous_names.append(name)

        upcoming_names = list()
        upcoming_person = max(upcoming_persons, key=operator.itemgetter(1))
        upcoming_delta = upcoming_person[1]
        upcoming_delta_days = abs(upcoming_delta.days)
        for name, _ in __same_by_delta(upcoming_persons, upcoming_delta):
            upcoming_names.append(name)

        out = list()

        if today_names:
            substitutes = {
                "name": humanize_names_list(today_names),
                "interjection": random.choice(INTERJECTIONS),
            }
            out.append(TODAY.format(**substitutes))

        def __day_or_days(number):
            if number > 1:
                return DAY_OR_DAYS["plural"]
            return DAY_OR_DAYS["singular"]

        substitutes = {
            "name": humanize_names_list(previous_names),
            "delta_days": previous_delta_days,
            "day_or_days": __day_or_days(previous_delta_days),
        }
        out.append(PREVIOUS.format(**substitutes))

        substitutes = {
            "name": humanize_names_list(upcoming_names),
            "delta_days": upcoming_delta_days,
            "day_or_days": __day_or_days(upcoming_delta_days),
        }
        out.append(UPCOMING.format(**substitutes))

        self.output.append(ChatMessage(message.Chat.Name, "\n".join(out)))

        return message, status

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Updated doctests<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
:mod:`birthdayreminder` --- Birthday reminder plugin
====================================================
"""

from __future__ import unicode_literals

__docformat__ = "restructuredtext en"

import atexit
import datetime
from threading import Timer
import random
import operator

from Skype4Py.enums import cmsReceived

from plugin import Plugin
from output import ChatMessage

INTERJECTIONS = [
    "Hurrah",
    "Hooray",
    "Huzzah",
    "Fancy",
    "Awesome",
    "Whee",
    "Whoa",
    "Unbelievable",
    "Wow",
    "Amazing",
    "Yahoo",
    "Yippie",
    "Yay",
    "Yeah",
    "Exciting",
    "Yee-haw",
    "Sweet",
    "Fantastic",
    "OMG",
]

TODAY = "Today is {name} birthday! {interjection}!"

PREVIOUS = "It was {name} birthday {delta_days} {day_or_days} ago."

UPCOMING = "{name} birthday is coming up next in {delta_days} {day_or_days}."

DAY_OR_DAYS = {
    "singular": "day",
    "plural": "days",
}

POSSESSIVE = "{name}'s"

CONJUNCTION = "and"

def humanize_names_list(names_list,
                        possessive=POSSESSIVE,
                        conjunction=CONJUNCTION):
    """
    >>> possessive = "{name}'s"
    >>> conjunction = "and"

    >>> humanize_names_list(["herp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's"

    >>> humanize_names_list(["herp", "derp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's and derp's"

    >>> humanize_names_list(["herp", "derp", "durp"],
    ...                     possessive=possessive,
    ...                     conjunction=conjunction)
    u"herp's, derp's and durp's"
    """

    assert isinstance(names_list, (list, tuple))
    names = [possessive.format(name=name) for name in names_list[:]]
    if len(names) is 1:
        return names[0]
    elif len(names) > 1:
        return "{0} {conjunction} {1}".format(", ".join(names[:-1]), names[-1],
                                              conjunction=conjunction)

def str_to_datetime(date_str):
    """
    >>> str_to_datetime("13.1")
    datetime.datetime(1900, 1, 13, 0, 0)

    >>> str_to_datetime("13.01")
    datetime.datetime(1900, 1, 13, 0, 0)

    >>> str_to_datetime("01.13")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format

    >>> str_to_datetime("13.1.2000")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("2000-01-13")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("2000-13-01")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format

    >>> str_to_datetime("13.01.2000")
    datetime.datetime(2000, 1, 13, 0, 0)

    >>> str_to_datetime("01.13.2000")
    Traceback (most recent call last):
        ...
    ValueError: Input string does not match any known date format
    """

    dt = None
    formats = ["%d.%m", "%d.%m.%Y", "%Y-%m-%d"]
    for fmt in formats:
        try:
            dt = datetime.datetime.strptime(date_str, fmt)
            break
        except ValueError:
            pass
    if dt is None:
        raise ValueError("Input string does not match any known date format")
    return dt

_timer = None

class BirthdayReminder(Plugin):
    """
    >>> import datetime
    >>> class DateMock(datetime.date):
    ...     @classmethod
    ...     def today(cls):
    ...         return datetime.datetime(2014, 3, 27)
    >>> datetime.date = DateMock
    >>> from Skype4Py.enums import cmsReceived
    >>> class ChatStub:
    ...     Name = "test"
    >>> class MessageStub:
    ...     Body = "!birthdays"
    ...     Chat = ChatStub
    >>> fake_status = cmsReceived

    >>> birthdays = {
    ...     "test1": "1986-03-27",
    ...     "test2": "2000-03-27",
    ...     "test3": "1986-03-25",
    ...     "test4": "2000-04-30",
    ... }
    >>> plugin = BirthdayReminder(birthdays=birthdays)
    >>> message, status = plugin.on_message_status(MessageStub, fake_status)
    >>> output = plugin.flush_output()[0]
    >>> output.text # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
    u"Today is test1's and test2's...test3's...2 days...test4's...34 days."
    """

    # Time in seconds.
    CHECK_INTERVAL = 600

    @staticmethod
    @atexit.register
    def _cleanup():
        global _timer
        try:
            _timer.cancel()
            _timer.join()
        except AttributeError:
            pass
        finally:
            del _timer

    def __init__(self, priority=0, whitelist=None, **kwargs):
        super(BirthdayReminder, self).__init__(priority, whitelist, **kwargs)

        self.dates = dict()
        birthdays = self.options.setdefault("birthdays")
        for name, date_str in birthdays.iteritems():
            try:
                dt = str_to_datetime(date_str)
            except ValueError:
                self._logger.error("Parse error: '%s' for '%s'", name, date_str)
            else:
                self.dates.update({name: dt})

        if not self.dates:
            self._logger.warning("No valid dates has been parsed!")
            return

        self._notified = False
        self._logger.info("Parsed %d date(s)", len(self.dates))
        self._check_dates()

    def _check_dates(self):
        global _timer

        today = datetime.date.today()

        persons = list()
        for name, dt in self.dates.iteritems():
            dt = dt.replace(year=today.year)
            delta = today - dt
            if delta.days == 0:
                persons.append(name)
        if persons:
            if not self._notified:
                self._notify_chats(persons)
                self._notified = True
            return
        self._notified = False
        _timer = Timer(self.CHECK_INTERVAL, self._check_dates)
        _timer.daemon = True
        _timer.start()

    def _notify_chats(self, names):
        substitutes = {
            "name": humanize_names_list(names),
            "interjection": random.choice(INTERJECTIONS),
        }
        message = TODAY.format(**substitutes)

        if self.whitelist:
            for chat in self.whitelist:
                self.output.append(ChatMessage(chat, message))

    def on_message_status(self, message, status):
        if status != cmsReceived:
            return

        if not message.Body.strip().startswith("!birthdays"):
            return

        today = datetime.date.today()

        previous_persons = list()
        upcoming_persons = list()
        today_names = list()
        for name, dt in self.dates.iteritems():
            dt = dt.replace(year=today.year)
            delta = today - dt
            person = (name, delta)
            if delta.days > 0:
                previous_persons.append(person)
            elif delta.days < 0:
                upcoming_persons.append(person)
            elif delta.days == 0:
                today_names.append(name)

        def __same_by_delta(persons, delta):
            for _name, _delta in persons:
                if _delta.days == delta.days:
                    yield (_name, _delta)

        previous_names = list()
        previous_person = min(previous_persons, key=operator.itemgetter(1))
        previous_delta = previous_person[1]
        previous_delta_days = previous_delta.days
        for name, _ in __same_by_delta(previous_persons, previous_delta):
            previous_names.append(name)

        upcoming_names = list()
        upcoming_person = max(upcoming_persons, key=operator.itemgetter(1))
        upcoming_delta = upcoming_person[1]
        upcoming_delta_days = abs(upcoming_delta.days)
        for name, _ in __same_by_delta(upcoming_persons, upcoming_delta):
            upcoming_names.append(name)

        out = list()

        if today_names:
            substitutes = {
                "name": humanize_names_list(today_names),
                "interjection": random.choice(INTERJECTIONS),
            }
            out.append(TODAY.format(**substitutes))

        def __day_or_days(number):
            if number > 1:
                return DAY_OR_DAYS["plural"]
            return DAY_OR_DAYS["singular"]

        substitutes = {
            "name": humanize_names_list(previous_names),
            "delta_days": previous_delta_days,
            "day_or_days": __day_or_days(previous_delta_days),
        }
        out.append(PREVIOUS.format(**substitutes))

        substitutes = {
            "name": humanize_names_list(upcoming_names),
            "delta_days": upcoming_delta_days,
            "day_or_days": __day_or_days(upcoming_delta_days),
        }
        out.append(UPCOMING.format(**substitutes))

        self.output.append(ChatMessage(message.Chat.Name, "\n".join(out)))

        return message, status

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


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

"""
:mod:`youtubeurlparser` --- YouTube URL parser plugin
=====================================================
"""

from __future__ import unicode_literals

__docformat__ = "restructuredtext en"

import re
import urllib2
import urllib
import urlparse
import json

from Skype4Py.enums import cmsReceived, cmsSent

from plugin import Plugin
from output import ChatMessage
import config

def get_video_id(url):
    """
    Parses video URL and returns its ID.

    :param url: YouTube video URL
    :type url: `str`

    :returns: YouTube video ID
    :rtype: `unicode` or `None`

    >>> get_video_id("https://www.youtube.com/watch?v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("https://www.youtube.com/watch?v=XXXXXXXXXXXZ")
    u'XXXXXXXXXXX'
    >>> get_video_id("https://www.youtube.com///watch?v=XXXXXXXXXXX///")
    u'XXXXXXXXXXX'
    >>> get_video_id("youtube.com/watch?v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be/XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://www.youtube.com/watch?feature=&v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be/YYYYYYYYYY") is None
    True
    >>> get_video_id("https://www.youtube.com/watch?v=") is None
    True
    >>> get_video_id("https://www.youtube.com/watch?v=X") is None
    True
    >>> get_video_id("http://youtu.be/XXXXXXXXXXX?t=0m00s")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be///X///") is None
    True
    >>> get_video_id("http://youtu.be///XXXXXXXXXXX///")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be///XXXXXXXXXXXZ///")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://www.youtube.com") is None
    True
    >>> get_video_id("youtube.com/watch?v=XXXXXXXXXXX&feature=youtu.be")
    u'XXXXXXXXXXX'
    >>> get_video_id("youtube.com/watch?feature=youtu.be&v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    """

    vid_id = ""

    if "youtube.com" in url:
        queries = urlparse.parse_qs(urlparse.urlparse(url).query)
        v = queries.get("v")
        try:
            vid_id = v[0].strip("/")[:11]
        except (IndexError, TypeError):
            return

    elif "youtu.be" in url:
        vid_id = [chunk.strip() for chunk in url.split("/") if chunk][-1][:11]

    return vid_id if len(vid_id) is 11 else None

def get_duration(youtube_duration):
    """
    >>> get_duration('PT15M51S')
    951

    >>> get_duration('PT1H11M10S')
    4270

    :param unicode youtube_duration: YouTube weirdly formatted duration
    :return int: duration in seconds for the human beings!
    """

    youtube_duration = youtube_duration.lower()

    duration = 0

    pattern = re.compile(r'([0-9]+)([HMS])', re.I | re.U)

    matches = re.finditer(pattern, youtube_duration)
    if not matches:
        return duration

    for part in matches:
        value, letter = part.groups()
        value = int(value)
        if letter == 'h':
            duration += value * 3600
        elif letter == 'm':
            duration += value * 60
        elif letter == 's':
            duration += value
        else:
            raise Exception("Unsupported letter: %s" % letter)

    return duration

class YouTubeURLParser(Plugin):
    """This plugin monitors received messages and outputs video title and its
    duration if that message contains a valid YouTube video URL.
    """

    _args = {
        'part': 'snippet,contentDetails',
        'fields': 'items(snippet(title),contentDetails(duration))',
        'key': config.GOOGLE_API_KEY,
    }

    _api_url = 'https://www.googleapis.com/youtube/v3/videos'

    _pattern = re.compile(ur"((?:youtube\.com|youtu\.be)/\S+)")

    _headers = {
        "User-Agent": "Gooby",
        "Accept-Language": "en-US,en;q=0.5",
        "Connection": "Keep-Alive",
    }
    _opener = urllib2.build_opener()
    _opener.addheaders = [(k, v) for k, v in _headers.iteritems()]

    def get_video_title(self, video_id):
        """Retrieves YouTube video title by its ID.

        :param video_id: YouTube video ID
        :type video_id: `str` or None

        >>> plugin = YouTubeURLParser()
        >>> plugin.get_video_title("dQw4w9WgXcQ")
        u'Rick Astley - Never Gonna Give You Up [00:03:33]'
        """

        cached_title = self._cache.get(video_id)
        if cached_title is not None:
            return cached_title

        args = self._args.copy()
        args.update({'id': video_id})
        url = ''.join((self._api_url, '?', urllib.urlencode(args)))

        try:
            data = json.loads(self._opener.open(url).read())
        except urllib2.HTTPError as error:
            raise Exception("Check your Google API key")

        try:
            title = data['items'][0]['snippet']['title']
        except (IndexError, AttributeError):
            title = ''

        try:
            duration = data['items'][0]['contentDetails']['duration']
        except (IndexError, AttributeError):
            duration = ''

        duration = get_duration(duration)

        minutes, seconds = divmod(int(duration), 60)
        hours, minutes = divmod(minutes, 60)
        duration = "{0:02d}:{1:02d}:{2:02d}".format(hours, minutes, seconds)

        title = u"{0} [{1}]".format(title, duration)

        self._cache.set(video_id, title)
        return title

    def on_message_status(self, message, status):
        if status not in (cmsReceived, cmsSent):
            return

        if not any(s in message.Body for s in ("youtu.be", "youtube.com")):
            return

        found = re.findall(self._pattern, message.Body)
        if not found:
            return

        titles = []

        for url in found:
            video_id = get_video_id(url)

            if video_id is None:
                continue

            self._logger.info("Retrieving {0} for {1}".format(
                video_id, message.FromHandle
            ))

            title = self.get_video_title(video_id)

            if title is not None:
                titles.append(title)
            else:
                msg = "Unable to retrieve video title for {0}".format(video_id)
                titles.append(msg)

                msg = "Unable to retrieve {0} for {1}".format(
                    video_id, message.FromHandle
                )
                self._logger.error(msg)

        if not titles:
            return

        if len(titles) is 1:
            msg = u"[YouTube] {0}".format("".join(titles))
        else:
            msg = u"[YouTube]\n{0}".format("\n".join(titles))
        #message.Chat.SendMessage(msg)
        self.output.append(ChatMessage(message.Chat.Name, msg))

        return message, status

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Minor fixes<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
:mod:`youtubeurlparser` --- YouTube URL parser plugin
=====================================================
"""

from __future__ import unicode_literals

__docformat__ = "restructuredtext en"

import re
import urllib2
import urllib
import urlparse
import json

from Skype4Py.enums import cmsReceived, cmsSent

from plugin import Plugin
from output import ChatMessage
import config

def get_video_id(url):
    """
    Parses video URL and returns its ID.

    :param url: YouTube video URL
    :type url: `str`

    :returns: YouTube video ID
    :rtype: `unicode` or `None`

    >>> get_video_id("https://www.youtube.com/watch?v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("https://www.youtube.com/watch?v=XXXXXXXXXXXZ")
    u'XXXXXXXXXXX'
    >>> get_video_id("https://www.youtube.com///watch?v=XXXXXXXXXXX///")
    u'XXXXXXXXXXX'
    >>> get_video_id("youtube.com/watch?v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be/XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://www.youtube.com/watch?feature=&v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be/YYYYYYYYYY") is None
    True
    >>> get_video_id("https://www.youtube.com/watch?v=") is None
    True
    >>> get_video_id("https://www.youtube.com/watch?v=X") is None
    True
    >>> get_video_id("http://youtu.be/XXXXXXXXXXX?t=0m00s")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be///X///") is None
    True
    >>> get_video_id("http://youtu.be///XXXXXXXXXXX///")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://youtu.be///XXXXXXXXXXXZ///")
    u'XXXXXXXXXXX'
    >>> get_video_id("http://www.youtube.com") is None
    True
    >>> get_video_id("youtube.com/watch?v=XXXXXXXXXXX&feature=youtu.be")
    u'XXXXXXXXXXX'
    >>> get_video_id("youtube.com/watch?feature=youtu.be&v=XXXXXXXXXXX")
    u'XXXXXXXXXXX'
    """

    vid_id = ""

    if "youtube.com" in url:
        queries = urlparse.parse_qs(urlparse.urlparse(url).query)
        v = queries.get("v")
        try:
            vid_id = v[0].strip("/")[:11]
        except (IndexError, TypeError):
            return

    elif "youtu.be" in url:
        vid_id = [chunk.strip() for chunk in url.split("/") if chunk][-1][:11]

    return vid_id if len(vid_id) is 11 else None

def get_duration(youtube_duration):
    """
    >>> get_duration('PT15M51S')
    951

    >>> get_duration('PT1H11M10S')
    4270

    :param unicode youtube_duration: YouTube weirdly formatted duration
    :return int: duration in seconds for the human beings!
    """

    youtube_duration = youtube_duration.lower()

    duration = 0

    pattern = re.compile(r'([0-9]+)([hms])', re.I | re.U)

    matches = re.finditer(pattern, youtube_duration)
    if not matches:
        return duration

    for match in matches:
        value, letter = match.groups()
        value = int(value)
        if letter == 'h':
            duration += value * 3600
        elif letter == 'm':
            duration += value * 60
        elif letter == 's':
            duration += value
        else:
            raise Exception("Unsupported letter: %s" % letter)

    return duration

class YouTubeURLParser(Plugin):
    """This plugin monitors received messages and outputs video title and its
    duration if that message contains a valid YouTube video URL.
    """

    _url_args = {
        'part': 'snippet,contentDetails',
        'fields': 'items(snippet(title),contentDetails(duration))',
        'key': config.GOOGLE_API_KEY,
    }

    _api_base_url = 'https://www.googleapis.com/youtube/v3/videos'

    _pattern = re.compile(ur"((?:youtube\.com|youtu\.be)/\S+)")

    _headers = {
        "User-Agent": "Gooby",
        "Accept-Language": "en-US,en;q=0.5",
        "Connection": "Keep-Alive",
    }
    _opener = urllib2.build_opener()
    _opener.addheaders = [(k, v) for k, v in _headers.iteritems()]

    def get_video_title(self, video_id):
        """Retrieves YouTube video title by its ID.

        :param video_id: YouTube video ID
        :type video_id: `str` or None

        >>> plugin = YouTubeURLParser()
        >>> plugin.get_video_title("dQw4w9WgXcQ")
        u'Rick Astley - Never Gonna Give You Up [00:03:33]'
        """

        cached_title = self._cache.get(video_id)
        if cached_title is not None:
            return cached_title

        args = self._url_args.copy()
        args.update({'id': video_id})
        url = ''.join((self._api_base_url, '?', urllib.urlencode(args)))

        try:
            data = json.loads(self._opener.open(url).read())
        except urllib2.HTTPError:
            raise Exception("Check your Google API key")

        try:
            title = data['items'][0]['snippet']['title']
        except (IndexError, AttributeError):
            title = None

        try:
            duration = data['items'][0]['contentDetails']['duration']
        except (IndexError, AttributeError):
            duration = None

        if None in (title, duration):
            raise Exception("Unable to retrieve video title parts")

        duration = get_duration(duration)

        minutes, seconds = divmod(int(duration), 60)
        hours, minutes = divmod(minutes, 60)
        duration = "{0:02d}:{1:02d}:{2:02d}".format(hours, minutes, seconds)

        title = u"{0} [{1}]".format(title, duration)

        self._cache.set(video_id, title)
        return title

    def on_message_status(self, message, status):
        if status not in (cmsReceived, cmsSent):
            return

        if not any(s in message.Body for s in ("youtu.be", "youtube.com")):
            return

        found = re.findall(self._pattern, message.Body)
        if not found:
            return

        titles = []

        for url in found:
            video_id = get_video_id(url)

            if video_id is None:
                continue

            self._logger.info("Retrieving {0} for {1}".format(
                video_id, message.FromHandle
            ))

            title = self.get_video_title(video_id)

            if title is not None:
                titles.append(title)
            else:
                msg = "Unable to retrieve video title for {0}".format(video_id)
                titles.append(msg)

                msg = "Unable to retrieve {0} for {1}".format(
                    video_id, message.FromHandle
                )
                self._logger.error(msg)

        if not titles:
            return

        if len(titles) is 1:
            msg = u"[YouTube] {0}".format("".join(titles))
        else:
            msg = u"[YouTube]\n{0}".format("\n".join(titles))
        #message.Chat.SendMessage(msg)
        self.output.append(ChatMessage(message.Chat.Name, msg))

        return message, status

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>import unittest
from functools import partial

from tvtk.api import tvtk
from simphony.cuds.particles import Particle, Bond, Particles
from simphony.core.data_container import DataContainer
from simphony.core.cuba import CUBA
from simphony.testing.utils import (
    compare_data_containers, compare_particles, compare_bonds)
from simphony.testing.abc_check_particles import (
    ContainerManipulatingBondsCheck, ContainerAddParticlesCheck,
    ContainerAddBondsCheck, ContainerManipulatingParticlesCheck)

from simphony_mayavi.cuds.api import VTKParticles
from simphony_mayavi.core.api import (
    supported_cuba, VTKEDGETYPES, VTKFACETYPES)

class TestVTKParticlesParticleOperations(
        ContainerAddParticlesCheck, unittest.TestCase):

    def supported_cuba(self):
        return supported_cuba()

    def container_factory(self, name):
        return VTKParticles(name=name)

class TestVTKParticlesManipulatingParticles(
        ContainerManipulatingParticlesCheck, unittest.TestCase):

    def supported_cuba(self):
        return supported_cuba()

    def container_factory(self, name):
        return VTKParticles(name=name)

class TestVTKParticlesAddBonds(
        ContainerAddBondsCheck, unittest.TestCase):

    def container_factory(self, name):
        return VTKParticles(name=name)

    def supported_cuba(self):
        return supported_cuba()

class TestVTKParticlesManipulatingBonds(
        ContainerManipulatingBondsCheck, unittest.TestCase):

    def container_factory(self, name):
        return VTKParticles(name=name)

    def supported_cuba(self):
        return supported_cuba()

class TestVTKParticlesDataContainer(unittest.TestCase):

    def setUp(self):
        self.addTypeEqualityFunc(
            DataContainer, partial(compare_data_containers, testcase=self))
        self.addTypeEqualityFunc(
            Particle, partial(compare_particles, testcase=self))
        self.addTypeEqualityFunc(
            Bond, partial(compare_bonds, testcase=self))

    def test_data(self):
        container = VTKParticles(name='foo')
        data = container.data
        data[CUBA.MASS] = 9
        container.data = data
        ret_data = container.data
        self.assertEqual(data, ret_data)
        self.assertIsNot(data, ret_data)
        cur_data = container.data
        cur_data[CUBA.MASS] = 10
        ret_data = container.data
        self.assertEqual(data, ret_data)
        self.assertIsNot(data, ret_data)

    def test_initialization_with_expected_size(self):
        data_set = tvtk.PolyData(points=tvtk.Points(), lines=[])
        container = VTKParticles(name='test', data_set=data_set)
        particle = Particle(coordinates=(0.0, 1.0, 2.0), data=DataContainer())
        uid = container.add_particle(particle)
        self.assertTrue(container.has_particle(uid))
        self.assertEqual(list(container.iter_bonds()), [])
        self.assertEqual(len(tuple(container.iter_particles())), 1)
        self.assertEqual(container.get_particle(uid), particle)

    def test_initialization_with_cuds(self):
        # given
        points = [
            [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
        bonds = [[0, 1], [0, 3], [1, 3, 2]]
        point_temperature = [10., 20., 30., 40.]
        bond_temperature = [60., 80., 190., 5.]
        reference = Particles('test')
        point_uids = [
            reference.add_particle(
                Particle(
                    coordinates=point,
                    data=DataContainer(
                        TEMPERATURE=point_temperature[index])))
            for index, point in enumerate(points)]
        bond_uids = [
            reference.add_bond(
                Bond(
                    particles=[point_uids[index] for index in indices],
                    data=DataContainer(
                        TEMPERATURE=bond_temperature[index])))
            for index, indices in enumerate(bonds)]

        # when
        container = VTKParticles.from_particles(reference)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, len(point_uids))
        for expected in reference.iter_particles():
            self.assertEqual(container.get_particle(expected.uid), expected)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, len(bond_uids))
        for expected in reference.iter_bonds():
            self.assertEqual(container.get_bond(expected.uid), expected)

    def test_initialization_with_empty_cuds(self):
        # given
        reference = Particles('test')

        # when
        container = VTKParticles.from_particles(reference)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 0)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 0)

    def test_initialization_with_poly_data(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.PolyData(points=points, lines=bonds)

        # when
        container = VTKParticles.from_dataset(name='test', data_set=vtk)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 12)
        for particle in container.iter_particles():
            self.assertIsNotNone(particle.uid)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 5)
        uids = [particle.uid for particle in container.iter_particles()]
        for bond in container.iter_bonds():
            self.assertIsNotNone(bond.uid)
            self.assertTrue(set(bond.particles).issubset(uids))

    def test_initialization_with_unstructured_grid(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.UnstructuredGrid(points=points)
        for bond in bonds:
            if len(bond) > 2:
                vtk.insert_next_cell(VTKEDGETYPES[1], bond)
            else:
                vtk.insert_next_cell(VTKEDGETYPES[0], bond)

        # when
        container = VTKParticles.from_dataset(name='test', data_set=vtk)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 12)
        for particle in container.iter_particles():
            self.assertIsNotNone(particle.uid)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 5)
        uids = [particle.uid for particle in container.iter_particles()]
        for bond in container.iter_bonds():
            self.assertIsNotNone(bond.uid)
            self.assertTrue(set(bond.particles).issubset(uids))

    def test_initialization_with_invalid_unstructured_grid(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.UnstructuredGrid(points=points)
        for bond in bonds:
            if len(bond) > 2:
                vtk.insert_next_cell(VTKEDGETYPES[1], bond)
            else:
                vtk.insert_next_cell(VTKEDGETYPES[0], bond)
        vtk.insert_next_cell(VTKFACETYPES[0], [0, 1, 2])

        # when/then
        with self.assertRaises(TypeError):
            VTKParticles.from_dataset(name='test', data_set=vtk)

    def test_initialization_with_image_data(self):
        # given
        vtk = tvtk.ImageData()
        vtk.extent = 0, 4, 0, 2, 0, 14

        # when/then
        with self.assertRaises(TypeError):
            VTKParticles.from_dataset(name='test', data_set=vtk)

if __name__ == '__main__':
    unittest.main()
<commit_msg>updated tests for vtk-particle<commit_after>import unittest
from functools import partial
import random

from tvtk.api import tvtk
from simphony.cuds.particles import Particle, Bond, Particles
from simphony.core.data_container import DataContainer
from simphony.core.cuba import CUBA
from simphony.testing.utils import (
    compare_data_containers, compare_particles, compare_bonds,
    create_bonds_with_id, create_data_container)
from simphony.testing.abc_check_particles import (
    CheckParticlesContainer,
    CheckAddingParticles, CheckManipulatingParticles,
    CheckAddingBonds, CheckManipulatingBonds)

from simphony_mayavi.cuds.api import VTKParticles
from simphony_mayavi.core.api import (
    supported_cuba, VTKEDGETYPES, VTKFACETYPES)

class TestVTKParticlesAddingParticlesOperations(
        CheckAddingParticles, unittest.TestCase):

    def supported_cuba(self):
        return supported_cuba()

    def container_factory(self, name):
        return VTKParticles(name=name)

class TestVTKParticlesManipulatingParticles(
        CheckManipulatingParticles, unittest.TestCase):

    def supported_cuba(self):
        return supported_cuba()

    def container_factory(self, name):
        return VTKParticles(name=name)

class TestVTKParticlesAddingBonds(
        CheckAddingBonds, unittest.TestCase):

    def container_factory(self, name):
        return VTKParticles(name=name)

    def supported_cuba(self):
        return supported_cuba()

class TestVTKParticlesManipulatingBonds(
        CheckManipulatingBonds, unittest.TestCase):

    def container_factory(self, name):
        return VTKParticles(name=name)

    def supported_cuba(self):
        return supported_cuba()

class TestVTKParticlesContainer(CheckParticlesContainer, unittest.TestCase):

    def container_factory(self, name):
        return VTKParticles(name=name)

    def supported_cuba(self):
        return supported_cuba()

class TestVTKParticlesDataContainer(unittest.TestCase):

    def setUp(self):
        self.addTypeEqualityFunc(
            DataContainer, partial(compare_data_containers, testcase=self))
        self.addTypeEqualityFunc(
            Particle, partial(compare_particles, testcase=self))
        self.addTypeEqualityFunc(
            Bond, partial(compare_bonds, testcase=self))

    def test_data(self):
        container = VTKParticles(name='foo')
        data = container.data
        data[CUBA.MASS] = 9
        container.data = data
        ret_data = container.data
        self.assertEqual(data, ret_data)
        self.assertIsNot(data, ret_data)
        cur_data = container.data
        cur_data[CUBA.MASS] = 10
        ret_data = container.data
        self.assertEqual(data, ret_data)
        self.assertIsNot(data, ret_data)

    def test_initialization_with_expected_size(self):
        data_set = tvtk.PolyData(points=tvtk.Points(), lines=[])
        container = VTKParticles(name='test', data_set=data_set)
        particle = Particle(coordinates=(0.0, 1.0, 2.0), data=DataContainer())
        uid = container.add_particle(particle)
        self.assertTrue(container.has_particle(uid))
        self.assertEqual(list(container.iter_bonds()), [])
        self.assertEqual(len(tuple(container.iter_particles())), 1)
        self.assertEqual(container.get_particle(uid), particle)

    def test_initialization_with_cuds(self):
        # given
        points = [
            [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
        bonds = [[0, 1], [0, 3], [1, 3, 2]]
        point_temperature = [10., 20., 30., 40.]
        bond_temperature = [60., 80., 190., 5.]
        reference = Particles('test')
        point_uids = [
            reference.add_particle(
                Particle(
                    coordinates=point,
                    data=DataContainer(
                        TEMPERATURE=point_temperature[index])))
            for index, point in enumerate(points)]
        bond_uids = [
            reference.add_bond(
                Bond(
                    particles=[point_uids[index] for index in indices],
                    data=DataContainer(
                        TEMPERATURE=bond_temperature[index])))
            for index, indices in enumerate(bonds)]

        # when
        container = VTKParticles.from_particles(reference)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, len(point_uids))
        for expected in reference.iter_particles():
            self.assertEqual(container.get_particle(expected.uid), expected)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, len(bond_uids))
        for expected in reference.iter_bonds():
            self.assertEqual(container.get_bond(expected.uid), expected)

    def test_initialization_with_empty_cuds(self):
        # given
        reference = Particles('test')

        # when
        container = VTKParticles.from_particles(reference)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 0)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 0)

    def test_initialization_with_poly_data(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.PolyData(points=points, lines=bonds)

        # when
        container = VTKParticles.from_dataset(name='test', data_set=vtk)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 12)
        for particle in container.iter_particles():
            self.assertIsNotNone(particle.uid)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 5)
        uids = [particle.uid for particle in container.iter_particles()]
        for bond in container.iter_bonds():
            self.assertIsNotNone(bond.uid)
            self.assertTrue(set(bond.particles).issubset(uids))

    def test_initialization_with_unstructured_grid(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.UnstructuredGrid(points=points)
        for bond in bonds:
            if len(bond) > 2:
                vtk.insert_next_cell(VTKEDGETYPES[1], bond)
            else:
                vtk.insert_next_cell(VTKEDGETYPES[0], bond)

        # when
        container = VTKParticles.from_dataset(name='test', data_set=vtk)

        # then
        number_of_particles = sum(1 for _ in container.iter_particles())
        self.assertEqual(number_of_particles, 12)
        for particle in container.iter_particles():
            self.assertIsNotNone(particle.uid)
        number_of_bonds = sum(1 for _ in container.iter_bonds())
        self.assertEqual(number_of_bonds, 5)
        uids = [particle.uid for particle in container.iter_particles()]
        for bond in container.iter_bonds():
            self.assertIsNotNone(bond.uid)
            self.assertTrue(set(bond.particles).issubset(uids))

    def test_initialization_with_invalid_unstructured_grid(self):
        # given
        bonds = [
            [0, 1, 2, 3],
            [4, 5, 6, 7, 8, 9, 10, 11],
            [2, 7, 11],
            [1, 4],
            [1, 5, 8]]
        points = [(i, i*2, i*3) for i in range(12)]
        vtk = tvtk.UnstructuredGrid(points=points)
        for bond in bonds:
            if len(bond) > 2:
                vtk.insert_next_cell(VTKEDGETYPES[1], bond)
            else:
                vtk.insert_next_cell(VTKEDGETYPES[0], bond)
        vtk.insert_next_cell(VTKFACETYPES[0], [0, 1, 2])

        # when/then
        with self.assertRaises(TypeError):
            VTKParticles.from_dataset(name='test', data_set=vtk)

    def test_initialization_with_image_data(self):
        # given
        vtk = tvtk.ImageData()
        vtk.extent = 0, 4, 0, 2, 0, 14

        # when/then
        with self.assertRaises(TypeError):
            VTKParticles.from_dataset(name='test', data_set=vtk)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>import vtk
import time
import numpy as np
from ddapp.timercallback import TimerCallback

class OrbitController(TimerCallback):

    def __init__(self, view):
        TimerCallback.__init__(self)
        self.view = view
        self.orbitTime = 20.0

    def tick(self):
        speed = 360.0 / self.orbitTime
        degrees = self.elapsed * speed
        self.view.camera().Azimuth(degrees)
        self.view.render()

class Flyer(TimerCallback):

    def __init__(self, view):
        TimerCallback.__init__(self)
        self.view = view
        self.flyTime = 0.5
        self.startTime = 0.0
        self.maintainViewDirection = False
        self.positionZoom = 0.7

    def getCameraCopy(self):
        camera = vtk.vtkCamera()
        camera.DeepCopy(self.view.camera())
        return camera

    def zoomTo(self, newFocalPoint, newPosition=None):

        self.interp = vtk.vtkCameraInterpolator()
        self.interp.AddCamera(0.0, self.getCameraCopy())

        c = self.getCameraCopy()
        newFocalPoint = np.array(newFocalPoint)
        oldFocalPoint = np.array(c.GetFocalPoint())
        oldPosition = np.array(c.GetPosition())

        if newPosition is None:
            if self.maintainViewDirection:
                newPosition = oldPosition + (newFocalPoint - oldFocalPoint)
            else:
                newPosition = oldPosition
            newPosition += self.positionZoom*(newFocalPoint - newPosition)
            #newPosition = newFocalPoint - self.positionZoom*(newFocalPoint - newPosition)

        c.SetFocalPoint(newFocalPoint)
        c.SetPosition(newPosition)
        c.SetViewUp([0.0, 0.0, 1.0])

        self.interp.AddCamera(1.0, c)
        self.startTime = time.time()
        self.start()

    def tick(self):

        elapsed = time.time() - self.startTime
        t = elapsed / float(self.flyTime)

        self.interp.InterpolateCamera(t, self.view.camera())
        self.view.render()

        if t > 1.0:
            return False

class RobotModelFollower(object):

    def __init__(self, view, robotModel, jointController):
        self.view = view
        self.robotModel = robotModel
        self.jointController = jointController
        self.callbackId = None

    def start(self):
        self.callbackId = self.robotModel.connectModelChanged(self.onModelChanged)
        self.lastTrackPosition = np.array(self.jointController.q[:3])

    def stop(self):
        self.robotModel.disconnectModelChanged(self.callbackId)

    def getCameraCopy(self):
        camera = vtk.vtkCamera()
        camera.DeepCopy(self.view.camera())
        return camera

    def onModelChanged(self, model):
        newTrackPosition = np.array(self.jointController.q[:3])

        delta = newTrackPosition - self.lastTrackPosition
        #delta[0] = 0.0
        #delta[1] = 0.0
        #delta[2] = 0.0

        self.lastTrackPosition = newTrackPosition

        c = self.view.camera()

        oldFocalPoint = np.array(c.GetFocalPoint())
        oldPosition = np.array(c.GetPosition())

        c.SetFocalPoint(oldFocalPoint + delta)
        c.SetPosition(oldPosition + delta)

        self.view.render()
<commit_msg>fix divide by zero in cameracontrol.Flyer for zero fly time<commit_after>import vtk
import time
import numpy as np
from ddapp.timercallback import TimerCallback

class OrbitController(TimerCallback):

    def __init__(self, view):
        TimerCallback.__init__(self)
        self.view = view
        self.orbitTime = 20.0

    def tick(self):
        speed = 360.0 / self.orbitTime
        degrees = self.elapsed * speed
        self.view.camera().Azimuth(degrees)
        self.view.render()

class Flyer(TimerCallback):

    def __init__(self, view):
        TimerCallback.__init__(self)
        self.view = view
        self.flyTime = 0.5
        self.startTime = 0.0
        self.maintainViewDirection = False
        self.positionZoom = 0.7

    def getCameraCopy(self):
        camera = vtk.vtkCamera()
        camera.DeepCopy(self.view.camera())
        return camera

    def zoomTo(self, newFocalPoint, newPosition=None):

        self.interp = vtk.vtkCameraInterpolator()
        self.interp.AddCamera(0.0, self.getCameraCopy())

        c = self.getCameraCopy()
        newFocalPoint = np.array(newFocalPoint)
        oldFocalPoint = np.array(c.GetFocalPoint())
        oldPosition = np.array(c.GetPosition())

        if newPosition is None:
            if self.maintainViewDirection:
                newPosition = oldPosition + (newFocalPoint - oldFocalPoint)
            else:
                newPosition = oldPosition
            newPosition += self.positionZoom*(newFocalPoint - newPosition)
            #newPosition = newFocalPoint - self.positionZoom*(newFocalPoint - newPosition)

        c.SetFocalPoint(newFocalPoint)
        c.SetPosition(newPosition)
        c.SetViewUp([0.0, 0.0, 1.0])

        self.interp.AddCamera(1.0, c)
        self.startTime = time.time()
        self.start()

    def tick(self):

        elapsed = time.time() - self.startTime
        t = (elapsed / float(self.flyTime)) if self.flyTime > 0 else 1.0

        self.interp.InterpolateCamera(t, self.view.camera())
        self.view.render()

        if t >= 1.0:
            return False

class RobotModelFollower(object):

    def __init__(self, view, robotModel, jointController):
        self.view = view
        self.robotModel = robotModel
        self.jointController = jointController
        self.callbackId = None

    def start(self):
        self.callbackId = self.robotModel.connectModelChanged(self.onModelChanged)
        self.lastTrackPosition = np.array(self.jointController.q[:3])

    def stop(self):
        self.robotModel.disconnectModelChanged(self.callbackId)

    def getCameraCopy(self):
        camera = vtk.vtkCamera()
        camera.DeepCopy(self.view.camera())
        return camera

    def onModelChanged(self, model):
        newTrackPosition = np.array(self.jointController.q[:3])

        delta = newTrackPosition - self.lastTrackPosition
        #delta[0] = 0.0
        #delta[1] = 0.0
        #delta[2] = 0.0

        self.lastTrackPosition = newTrackPosition

        c = self.view.camera()

        oldFocalPoint = np.array(c.GetFocalPoint())
        oldPosition = np.array(c.GetPosition())

        c.SetFocalPoint(oldFocalPoint + delta)
        c.SetPosition(oldPosition + delta)

        self.view.render()
<|endoftext|>