<commit_before>from braces.views import CsrfExemptMixin, JsonRequestResponseMixin
from chartit import DataPool, Chart
from core.models import Player, Contract, PlayerStatistics, Finance, Matchday
from django.contrib.auth.decorators import login_required
from django.core.exceptions import MultipleObjectsReturned
from django.utils.decorators import method_decorator
from django.views.generic import DetailView, TemplateView, View

@method_decorator(login_required, name='dispatch')
class PlayerStatisticsView(TemplateView):
    template_name = 'core/ofm/player_statistics.html'

    def get_context_data(self, **kwargs):
        context = super(PlayerStatisticsView, self).get_context_data(**kwargs)

        matchdays = Matchday.objects.filter(player_statistics__isnull=False).distinct()

        context['matchdays'] = matchdays

        return context

@method_decorator(login_required, name='dispatch')
class PlayerStatisticsAsJsonView(CsrfExemptMixin, JsonRequestResponseMixin, View):

    def get(self, request, *args, **kwargs):
        contracts = Contract.objects.filter(user=self.request.user, sold_on_matchday=None)
        players = [contract.player for contract in contracts]

        newer_matchday_season = self.request.GET.get('newer_matchday_season', default=Matchday.objects.all()[0].season.number)
        newer_matchday = self.request.GET.get('newer_matchday', default=Matchday.objects.all()[0].number)
        older_matchday_season = self.request.GET.get('older_matchday_season', default=Matchday.objects.all()[1].season.number if Matchday.objects.all().count() > 1 else None)
        older_matchday = self.request.GET.get('older_matchday', default=Matchday.objects.all()[1].number if Matchday.objects.all().count() > 1 else None)

        show_diff = self.request.GET.get('show_diff', default='false').lower() == 'true'

        player_statistics_tuples = [self._get_statistics_from_player_and_matchday(player,
                                                                            newer_matchday_season, newer_matchday,
                                                                            older_matchday_season, older_matchday)
                             for player in players]

        player_statistics_json = [self._get_player_statistics_diff_in_json(newer_player_statistic,
                                                                           older_player_statistic, show_diff)
                                  for (newer_player_statistic, older_player_statistic) in player_statistics_tuples]

        return self.render_json_response(player_statistics_json)

    def _get_statistics_from_player_and_matchday(self, player,
                                                 newer_matchday_season, newer_matchday,
                                                 older_matchday_season, older_matchday):

        ps1 = PlayerStatistics.objects.filter(player=player, matchday__season__number=newer_matchday_season, matchday__number=newer_matchday)
        ps2 = PlayerStatistics.objects.filter(player=player, matchday__season__number=older_matchday_season, matchday__number=older_matchday)

        ps1 = self._validate_filtered_player_statistics(ps1)
        ps2 = self._validate_filtered_player_statistics(ps2)

        return ps1, ps2

    def _validate_filtered_player_statistics(self, player_statistics):
        if len(player_statistics) > 1:
            raise MultipleObjectsReturned
        elif player_statistics:
            player_statistics = player_statistics[0]
        return player_statistics

    def _get_player_statistics_diff_in_json(self, newer_player_statistics, older_player_statistics, show_diff=True):
        """
        Args:
            newer_player_statistics: newer statistic
            older_player_statistics: older statistic

        Returns:
            A dictionary of player statistics data. If st2 is None st1 is returned
        """

        if not newer_player_statistics:
            newer_player_statistics = PlayerStatistics.objects.all()[0]

        ep = newer_player_statistics.ep
        if show_diff and older_player_statistics:
            ep = newer_player_statistics.ep - older_player_statistics.ep
        tp = newer_player_statistics.tp
        if show_diff and older_player_statistics:
            tp = newer_player_statistics.tp - older_player_statistics.tp
        awp = newer_player_statistics.awp
        if show_diff and older_player_statistics:
            awp = newer_player_statistics.awp - older_player_statistics.awp
        freshness = newer_player_statistics.freshness
        if show_diff and older_player_statistics:
            freshness = newer_player_statistics.freshness - older_player_statistics.freshness

        statistic_diff = dict()
        statistic_diff['position'] = newer_player_statistics.player.position
        statistic_diff['age'] = newer_player_statistics.age
        statistic_diff['strength'] = newer_player_statistics.strength
        statistic_diff['name'] = '%s' % (newer_player_statistics.player.get_absolute_url(), newer_player_statistics.player.name)
        statistic_diff['ep'] = ep
        statistic_diff['tp'] = tp
        statistic_diff['awp'] = awp
        statistic_diff['freshness'] = freshness
        statistic_diff['games_in_season'] = newer_player_statistics.games_in_season
        statistic_diff['goals_in_season'] = newer_player_statistics.goals_in_season
        statistic_diff['won_tacklings_in_season'] = newer_player_statistics.won_tacklings_in_season
        statistic_diff['lost_tacklings_in_season'] = newer_player_statistics.lost_tacklings_in_season
        statistic_diff['won_friendly_tacklings_in_season'] = newer_player_statistics.won_friendly_tacklings_in_season
        statistic_diff['lost_friendly_tacklings_in_season'] = newer_player_statistics.lost_friendly_tacklings_in_season
        statistic_diff['yellow_cards_in_season'] = newer_player_statistics.yellow_cards_in_season
        statistic_diff['red_cards_in_season'] = newer_player_statistics.red_cards_in_season
        statistic_diff['equity'] = newer_player_statistics.equity

        return statistic_diff

@method_decorator(login_required, name='dispatch')
class PlayerDetailView(DetailView):
    context_object_name = 'player'
    template_name = 'core/ofm/player_detail.html'
    queryset = Player.objects.all()

    def get_context_data(self, **kwargs):
        context = super(PlayerDetailView, self).get_context_data(**kwargs)

        player = self.get_object()

        statistics_data = DataPool(
            series=[
                {'options':
                    {'source': PlayerStatistics.objects.filter(player=player)},
                    'terms': [
                        'matchday__number',
                        'ep',
                        'tp',
                        'awp'
                    ]
                }
            ]
        )

        chart = Chart(
            datasource=statistics_data,
            series_options=
            [{'options': {
                'type': 'spline',
                'stacking': False},
                'terms': {'matchday__number': ['ep', 'tp', 'awp', ]}
            }],
            chart_options=
            {
                'title': {
                    'text': 'Spielerstatistik'
                },
                'xAxis': {
                    'title': {
                       'text': 'Spieltag'}
                },
                'yAxis': {
                    'title': {
                       'text': ' '}
                },
            }
        )

        context['chart'] = chart

        return context

    def get_object(self):
        player = super(PlayerDetailView, self).get_object()
        contracts = Contract.objects.filter(user=self.request.user, player=player, sold_on_matchday=None)
        return player if contracts.count() > 0 else None

@method_decorator(login_required, name='dispatch')
class FinanceDataView(TemplateView):
    template_name = 'core/ofm/finances.html'

    def get_context_data(self, **kwargs):
        context = super(FinanceDataView, self).get_context_data(**kwargs)

        finance_data = DataPool(
            series=[
                {'options':
                    {'source': Finance.objects.filter(user=self.request.user)},
                    'terms': [
                        'matchday__number',
                        'balance',
                    ]
                }
            ]
        )

        chart = Chart(
            datasource=finance_data,
            series_options=
            [{'options': {
                'type': 'spline',
                'stacking': False},
                'terms': {'matchday__number': ['balance', ]}
            }],
            chart_options=
            {
                'title': {
                    'text': 'Finanzstatistik'
                },
                'xAxis': {
                    'title': {
                       'text': 'Spieltag'}
                },
                'yAxis': {
                    'title': {
                       'text': ' '}
                },
            }
        )

        context['chart'] = chart

        return context
<commit_msg>add new view for finance column chart<commit_after>from braces.views import CsrfExemptMixin, JsonRequestResponseMixin
from chartit import DataPool, Chart
from core.models import Player, Contract, PlayerStatistics, Finance, Matchday
from django.contrib.auth.decorators import login_required
from django.core.exceptions import MultipleObjectsReturned
from django.utils.decorators import method_decorator
from django.views.generic import DetailView, TemplateView, View

@method_decorator(login_required, name='dispatch')
class PlayerStatisticsView(TemplateView):
    template_name = 'core/ofm/player_statistics.html'

    def get_context_data(self, **kwargs):
        context = super(PlayerStatisticsView, self).get_context_data(**kwargs)

        matchdays = Matchday.objects.filter(player_statistics__isnull=False).distinct()

        context['matchdays'] = matchdays

        return context

@method_decorator(login_required, name='dispatch')
class PlayerStatisticsAsJsonView(CsrfExemptMixin, JsonRequestResponseMixin, View):

    def get(self, request, *args, **kwargs):
        contracts = Contract.objects.filter(user=self.request.user, sold_on_matchday=None)
        players = [contract.player for contract in contracts]

        newer_matchday_season = self.request.GET.get('newer_matchday_season', default=Matchday.objects.all()[0].season.number)
        newer_matchday = self.request.GET.get('newer_matchday', default=Matchday.objects.all()[0].number)
        older_matchday_season = self.request.GET.get('older_matchday_season', default=Matchday.objects.all()[1].season.number if Matchday.objects.all().count() > 1 else None)
        older_matchday = self.request.GET.get('older_matchday', default=Matchday.objects.all()[1].number if Matchday.objects.all().count() > 1 else None)

        show_diff = self.request.GET.get('show_diff', default='false').lower() == 'true'

        player_statistics_tuples = [self._get_statistics_from_player_and_matchday(player,
                                                                            newer_matchday_season, newer_matchday,
                                                                            older_matchday_season, older_matchday)
                             for player in players]

        player_statistics_json = [self._get_player_statistics_diff_in_json(newer_player_statistic,
                                                                           older_player_statistic, show_diff)
                                  for (newer_player_statistic, older_player_statistic) in player_statistics_tuples]

        return self.render_json_response(player_statistics_json)

    def _get_statistics_from_player_and_matchday(self, player,
                                                 newer_matchday_season, newer_matchday,
                                                 older_matchday_season, older_matchday):

        ps1 = PlayerStatistics.objects.filter(player=player, matchday__season__number=newer_matchday_season, matchday__number=newer_matchday)
        ps2 = PlayerStatistics.objects.filter(player=player, matchday__season__number=older_matchday_season, matchday__number=older_matchday)

        ps1 = self._validate_filtered_player_statistics(ps1)
        ps2 = self._validate_filtered_player_statistics(ps2)

        return ps1, ps2

    def _validate_filtered_player_statistics(self, player_statistics):
        if len(player_statistics) > 1:
            raise MultipleObjectsReturned
        elif player_statistics:
            player_statistics = player_statistics[0]
        return player_statistics

    def _get_player_statistics_diff_in_json(self, newer_player_statistics, older_player_statistics, show_diff=True):
        """
        Args:
            newer_player_statistics: newer statistic
            older_player_statistics: older statistic

        Returns:
            A dictionary of player statistics data. If st2 is None st1 is returned
        """

        if not newer_player_statistics:
            newer_player_statistics = PlayerStatistics.objects.all()[0]

        ep = newer_player_statistics.ep
        if show_diff and older_player_statistics:
            ep = newer_player_statistics.ep - older_player_statistics.ep
        tp = newer_player_statistics.tp
        if show_diff and older_player_statistics:
            tp = newer_player_statistics.tp - older_player_statistics.tp
        awp = newer_player_statistics.awp
        if show_diff and older_player_statistics:
            awp = newer_player_statistics.awp - older_player_statistics.awp
        freshness = newer_player_statistics.freshness
        if show_diff and older_player_statistics:
            freshness = newer_player_statistics.freshness - older_player_statistics.freshness

        statistic_diff = dict()
        statistic_diff['position'] = newer_player_statistics.player.position
        statistic_diff['age'] = newer_player_statistics.age
        statistic_diff['strength'] = newer_player_statistics.strength
        statistic_diff['name'] = '%s' % (newer_player_statistics.player.get_absolute_url(), newer_player_statistics.player.name)
        statistic_diff['ep'] = ep
        statistic_diff['tp'] = tp
        statistic_diff['awp'] = awp
        statistic_diff['freshness'] = freshness
        statistic_diff['games_in_season'] = newer_player_statistics.games_in_season
        statistic_diff['goals_in_season'] = newer_player_statistics.goals_in_season
        statistic_diff['won_tacklings_in_season'] = newer_player_statistics.won_tacklings_in_season
        statistic_diff['lost_tacklings_in_season'] = newer_player_statistics.lost_tacklings_in_season
        statistic_diff['won_friendly_tacklings_in_season'] = newer_player_statistics.won_friendly_tacklings_in_season
        statistic_diff['lost_friendly_tacklings_in_season'] = newer_player_statistics.lost_friendly_tacklings_in_season
        statistic_diff['yellow_cards_in_season'] = newer_player_statistics.yellow_cards_in_season
        statistic_diff['red_cards_in_season'] = newer_player_statistics.red_cards_in_season
        statistic_diff['equity'] = newer_player_statistics.equity

        return statistic_diff

@method_decorator(login_required, name='dispatch')
class PlayerDetailView(DetailView):
    context_object_name = 'player'
    template_name = 'core/ofm/player_detail.html'
    queryset = Player.objects.all()

    def get_context_data(self, **kwargs):
        context = super(PlayerDetailView, self).get_context_data(**kwargs)

        player = self.get_object()

        statistics_data = DataPool(
            series=[
                {'options':
                    {'source': PlayerStatistics.objects.filter(player=player)},
                    'terms': [
                        'matchday__number',
                        'ep',
                        'tp',
                        'awp'
                    ]
                }
            ]
        )

        chart = Chart(
            datasource=statistics_data,
            series_options=
            [{'options': {
                'type': 'spline',
                'stacking': False},
                'terms': {'matchday__number': ['ep', 'tp', 'awp', ]}
            }],
            chart_options=
            {
                'title': {
                    'text': 'Spielerstatistik'
                },
                'xAxis': {
                    'title': {
                       'text': 'Spieltag'}
                },
                'yAxis': {
                    'title': {
                       'text': ' '}
                },
            }
        )

        context['chart'] = chart

        return context

    def get_object(self):
        player = super(PlayerDetailView, self).get_object()
        contracts = Contract.objects.filter(user=self.request.user, player=player, sold_on_matchday=None)
        return player if contracts.count() > 0 else None

@method_decorator(login_required, name='dispatch')
class FinanceDataView(TemplateView):
    template_name = 'core/ofm/finances.html'

    def get_context_data(self, **kwargs):
        context = super(FinanceDataView, self).get_context_data(**kwargs)

        finance_data = DataPool(
            series=[
                {'options':
                    {'source': Finance.objects.filter(user=self.request.user)},
                    'terms': [
                        'matchday__number',
                        'balance',
                    ]
                }
            ]
        )

        chart = Chart(
            datasource=finance_data,
            series_options=
            [{'options': {
                    'type': 'spline',
                    'stacking': False,
                    'allowPointSelect': True,
                },
                'terms': {'matchday__number': ['balance', ]}
            }],
            chart_options=
            {
                'title': {
                    'text': 'Finanzstatistik'
                },
                'xAxis': {
                    'title': {
                       'text': 'Spieltag'}
                },
                'yAxis': {
                    'title': {
                       'text': ' '}
                },
            }
        )

        context['chart'] = chart

        return context

@method_decorator(login_required, name='dispatch')
class FinanceDataColumnChartView(TemplateView):
    template_name = 'core/ofm/finances.html'

    def get_context_data(self, **kwargs):
        context = super(FinanceDataColumnChartView, self).get_context_data(**kwargs)

        finance_data = DataPool(
            series=[
                {'options':
                    {'source': Finance.objects.filter(user=self.request.user)},
                    'terms': [
                        'matchday__number',
                        'balance',
                        'income_visitors_league',
                        'income_sponsoring',
                        'income_cup',
                        'income_interests',
                        'income_loan',
                        'income_transfer',
                        'income_visitors_friendlies',
                        'income_friendlies',
                        'income_funcup',
                        'income_betting',
                        'expenses_player_salaries',
                        'expenses_stadium',
                        'expenses_youth',
                        'expenses_interests',
                        'expenses_trainings',
                        'expenses_transfer',
                        'expenses_compensation',
                        'expenses_friendlies',
                        'expenses_funcup',
                        'expenses_betting',
                    ]
                }
            ]
        )

        chart = Chart(
            datasource=finance_data,
            series_options=
            [
                {'options': {
                        'type': 'column',
                        'stacking': True,
                        'stack': 0,
                    },
                    'terms': {'matchday__number': [
                                'income_visitors_league',
                                'income_sponsoring',
                                'income_cup',
                                'income_interests',
                                'income_loan',
                                'income_transfer',
                                'income_visitors_friendlies',
                                'income_friendlies',
                                'income_funcup',
                                'income_betting']}
                },
                {'options': {
                        'type': 'column',
                        'stacking': True,
                        'stack': 1,
                    },
                    'terms': {'matchday__number': [
                                'expenses_player_salaries',
                                'expenses_stadium',
                                'expenses_youth',
                                'expenses_interests',
                                'expenses_trainings',
                                'expenses_transfer',
                                'expenses_compensation',
                                'expenses_friendlies',
                                'expenses_funcup',
                                'expenses_betting']}
                }
            ],
            chart_options=
            {
                'title': {
                    'text': 'Finanzstatistik'
                },
                'xAxis': {
                    'title': {
                       'text': 'Spieltag'}
                },
                'yAxis': {
                    'title': {
                       'text': ' '}
                },
            }
        )

        context['chart'] = chart

        return context
<|endoftext|>


<commit_before>from couchdbkit.resource import ResourceNotFound
from dimagi.utils.logging import notify_exception

class SyncCouchToSQLMixin(object):
    """
    A mixin to be used for migrating a Couch model
    to a SQL model (see also SyncSQLToCouchMixin).

    Steps to use it:

    1. Create a Django ORM model which inherits from SyncSQLToCouchMixin
    and is identical to the Couch model with one additional field named
    couch_id, for the _id of the Couch doc. Make sure you index couch_id
    and allow it to be null since it can be initially null. Do not include
    this field in _migration_get_fields().

    2. Make the Couch model inherit from SyncCouchToSQLMixin.

    3. Implement the unimplemented methods from the mixins on both models.

    4. Make sure neither of the models overrides the save() or delete()
    methods. If either does, you will have to manually insert the sync
    code from the corresponding mixin. Also make sure the mixins are the
    left-most classes being inherited from.

    After that, any saves to the Couch model will sync the SQL model, and
    any saves to the SQL model will sync the Couch model. After migrating
    each doc, you can incrementally change the code to use to SQL model
    instead of the Couch model, and eventually remove the Couch model entirely.

    If you have a custom sync process, just override _migration_sync_to_sql.
    """

    @classmethod
    def _migration_get_fields(cls):
        """
        Should return a list of field names to sync. These field names should
        be identical between the Couch and SQL models.
        """
        raise NotImplementedError()

    @classmethod
    def _migration_get_sql_model_class(cls):
        """
        Should return the class of the SQL model.
        """
        raise NotImplementedError()

    def _migration_automatically_handle_dups(self):
        """
        Ideally, you should use a CriticalSection to lock out your get and save
        functionality so that each sync will automatically be atomic. However,
        even after doing so, during periods of high latency the locks may time
        out and you still might end up with a race condition in the sync that
        creates more than one SQL model from this couch model.

        Make this method return True if you want subsequent syncs to delete
        the duplicate records and resync. Otherwise, the default behavior is
        that a notify exception email will go out each time duplicates are found.
        """
        return False

    def _migration_get_sql_object(self):
        cls = self._migration_get_sql_model_class()
        try:
            return cls.objects.get(couch_id=self._id)
        except cls.DoesNotExist:
            return None
        except cls.MultipleObjectsReturned:
            if not self._migration_automatically_handle_dups():
                raise
            cls.objects.filter(couch_id=self._id).delete()
            return None

    def _migration_get_or_create_sql_object(self):
        cls = self._migration_get_sql_model_class()
        obj = self._migration_get_sql_object()
        if obj is None:
            obj = cls(couch_id=self._id)
        return obj

    def _migration_sync_to_sql(self, sql_object):
        for field_name in self._migration_get_fields():
            value = getattr(self, field_name)
            setattr(sql_object, field_name, value)
        sql_object.save(sync_to_couch=False)

    def _migration_do_sync(self):
        sql_object = self._migration_get_or_create_sql_object()
        self._migration_sync_to_sql(sql_object)

    def save(self, *args, **kwargs):
        sync_to_sql = kwargs.pop('sync_to_sql', True)
        super(SyncCouchToSQLMixin, self).save(*args, **kwargs)
        if sync_to_sql:
            try:
                self._migration_do_sync()
            except:
                sql_class_name = self._migration_get_sql_model_class().__name__
                couch_class_name = self.__class__.__name__
                notify_exception(None,
                    message='Could not sync %s SQL object from %s %s' % (sql_class_name,
                        couch_class_name, self._id))

    def delete(self, *args, **kwargs):
        sync_to_sql = kwargs.pop('sync_to_sql', True)
        if sync_to_sql:
            sql_object = self._migration_get_sql_object()
            if sql_object is not None:
                sql_object.delete(sync_to_couch=False)
        super(SyncCouchToSQLMixin, self).delete(*args, **kwargs)

class SyncSQLToCouchMixin(object):
    """
    See SyncCouchToSQLMixin.

    If you have a custom sync process, just override _migration_sync_to_couch.
    """

    @classmethod
    def _migration_get_fields(cls):
        """
        Should return a list of field names to sync. These field names should
        be identical between the Couch and SQL models.
        """
        raise NotImplementedError()

    @classmethod
    def _migration_get_couch_model_class(cls):
        """
        Should return the class of the Couch model.
        """
        raise NotImplementedError()

    def _migration_get_couch_object(self):
        if not self.couch_id:
            return None
        cls = self._migration_get_couch_model_class()
        return cls.get(self.couch_id)

    def _migration_get_or_create_couch_object(self):
        cls = self._migration_get_couch_model_class()
        obj = self._migration_get_couch_object()
        if obj is None:
            obj = cls()
            obj.save(sync_to_sql=False)
            self.couch_id = obj._id
            self.save(sync_to_couch=False)
        return obj

    def _migration_sync_to_couch(self, couch_obj):
        for field_name in self._migration_get_fields():
            value = getattr(self, field_name)
            setattr(couch_obj, field_name, value)
        couch_obj.save(sync_to_sql=False)

    def _migration_do_sync(self):
        couch_obj = self._migration_get_or_create_couch_object()
        self._migration_sync_to_couch(couch_obj)

    def save(self, *args, **kwargs):
        sync_to_couch = kwargs.pop('sync_to_couch', True)
        super(SyncSQLToCouchMixin, self).save(*args, **kwargs)
        if sync_to_couch:
            try:
                self._migration_do_sync()
            except:
                couch_class_name = self._migration_get_couch_model_class().__name__
                sql_class_name = self.__class__.__name__
                notify_exception(None,
                    message='Could not sync %s Couch object from %s %s' % (couch_class_name,
                        sql_class_name, self.pk))

    def delete(self, *args, **kwargs):
        sync_to_couch = kwargs.pop('sync_to_couch', True)
        if sync_to_couch:
            couch_object = self._migration_get_couch_object()
            if couch_object is not None:
                couch_object.delete(sync_to_sql=False)
        super(SyncSQLToCouchMixin, self).delete(*args, **kwargs)
<commit_msg>Allow setting of couch_id_name<commit_after>from dimagi.utils.logging import notify_exception

class SyncCouchToSQLMixin(object):
    """
    A mixin to be used for migrating a Couch model
    to a SQL model (see also SyncSQLToCouchMixin).

    Steps to use it:

    1. Create a Django ORM model which inherits from SyncSQLToCouchMixin
    and is identical to the Couch model with one additional field named
    couch_id, for the _id of the Couch doc. Make sure you index couch_id
    and allow it to be null since it can be initially null. Do not include
    this field in _migration_get_fields().
    You can override 'couch_id_name' if you need to use a different name.

    2. Make the Couch model inherit from SyncCouchToSQLMixin.

    3. Implement the unimplemented methods from the mixins on both models.

    4. Make sure neither of the models overrides the save() or delete()
    methods. If either does, you will have to manually insert the sync
    code from the corresponding mixin. Also make sure the mixins are the
    left-most classes being inherited from.

    After that, any saves to the Couch model will sync the SQL model, and
    any saves to the SQL model will sync the Couch model. After migrating
    each doc, you can incrementally change the code to use to SQL model
    instead of the Couch model, and eventually remove the Couch model entirely.

    If you have a custom sync process, just override _migration_sync_to_sql.
    """

    @classmethod
    def _migration_get_fields(cls):
        """
        Should return a list of field names to sync. These field names should
        be identical between the Couch and SQL models.
        """
        raise NotImplementedError()

    @classmethod
    def _migration_get_sql_model_class(cls):
        """
        Should return the class of the SQL model.
        """
        raise NotImplementedError()

    def _migration_automatically_handle_dups(self):
        """
        Ideally, you should use a CriticalSection to lock out your get and save
        functionality so that each sync will automatically be atomic. However,
        even after doing so, during periods of high latency the locks may time
        out and you still might end up with a race condition in the sync that
        creates more than one SQL model from this couch model.

        Make this method return True if you want subsequent syncs to delete
        the duplicate records and resync. Otherwise, the default behavior is
        that a notify exception email will go out each time duplicates are found.
        """
        return False

    def _migration_get_sql_object(self):
        cls = self._migration_get_sql_model_class()
        try:
            return cls.objects.get(**{cls.couch_id_name: self._id})
        except cls.DoesNotExist:
            return None
        except cls.MultipleObjectsReturned:
            if not self._migration_automatically_handle_dups():
                raise
            cls.objects.filter(**{cls.couch_id_name: self._id}).delete()
            return None

    def _migration_get_or_create_sql_object(self):
        cls = self._migration_get_sql_model_class()
        obj = self._migration_get_sql_object()
        if obj is None:
            obj = cls(**{cls.couch_id_name: self._id})
        return obj

    def _migration_sync_to_sql(self, sql_object):
        for field_name in self._migration_get_fields():
            value = getattr(self, field_name)
            setattr(sql_object, field_name, value)
        sql_object.save(sync_to_couch=False)

    def _migration_do_sync(self):
        sql_object = self._migration_get_or_create_sql_object()
        self._migration_sync_to_sql(sql_object)

    def save(self, *args, **kwargs):
        sync_to_sql = kwargs.pop('sync_to_sql', True)
        super(SyncCouchToSQLMixin, self).save(*args, **kwargs)
        if sync_to_sql:
            try:
                self._migration_do_sync()
            except:
                sql_class_name = self._migration_get_sql_model_class().__name__
                couch_class_name = self.__class__.__name__
                notify_exception(None,
                    message='Could not sync %s SQL object from %s %s' % (sql_class_name,
                        couch_class_name, self._id))

    def delete(self, *args, **kwargs):
        sync_to_sql = kwargs.pop('sync_to_sql', True)
        if sync_to_sql:
            sql_object = self._migration_get_sql_object()
            if sql_object is not None:
                sql_object.delete(sync_to_couch=False)
        super(SyncCouchToSQLMixin, self).delete(*args, **kwargs)

class SyncSQLToCouchMixin(object):
    """
    See SyncCouchToSQLMixin.

    If you have a custom sync process, just override _migration_sync_to_couch.
    """
    couch_id_name = "couch_id"

    @classmethod
    def _migration_get_fields(cls):
        """
        Should return a list of field names to sync. These field names should
        be identical between the Couch and SQL models.
        """
        raise NotImplementedError()

    @classmethod
    def _migration_get_couch_model_class(cls):
        """
        Should return the class of the Couch model.
        """
        raise NotImplementedError()

    def _migration_get_couch_object(self):
        if not getattr(self, self.couch_id_name, None):
            return None
        cls = self._migration_get_couch_model_class()
        return cls.get(getattr(self, self.couch_id_name, None))

    def _migration_get_or_create_couch_object(self):
        cls = self._migration_get_couch_model_class()
        obj = self._migration_get_couch_object()
        if obj is None:
            obj = cls()
            obj.save(sync_to_sql=False)
            setattr(self, self.couch_id_name, obj._id)
            self.save(sync_to_couch=False)
        return obj

    def _migration_sync_to_couch(self, couch_obj):
        for field_name in self._migration_get_fields():
            value = getattr(self, field_name)
            setattr(couch_obj, field_name, value)
        couch_obj.save(sync_to_sql=False)

    def _migration_do_sync(self):
        couch_obj = self._migration_get_or_create_couch_object()
        self._migration_sync_to_couch(couch_obj)

    def save(self, *args, **kwargs):
        sync_to_couch = kwargs.pop('sync_to_couch', True)
        super(SyncSQLToCouchMixin, self).save(*args, **kwargs)
        if sync_to_couch:
            try:
                self._migration_do_sync()
            except:
                couch_class_name = self._migration_get_couch_model_class().__name__
                sql_class_name = self.__class__.__name__
                notify_exception(None,
                    message='Could not sync %s Couch object from %s %s' % (couch_class_name,
                        sql_class_name, self.pk))

    def delete(self, *args, **kwargs):
        sync_to_couch = kwargs.pop('sync_to_couch', True)
        if sync_to_couch:
            couch_object = self._migration_get_couch_object()
            if couch_object is not None:
                couch_object.delete(sync_to_sql=False)
        super(SyncSQLToCouchMixin, self).delete(*args, **kwargs)
<|endoftext|>


<commit_before>import logging
import time
from datetime import datetime, date
from decimal import Decimal
from itertools import groupby

from django import db
from django.contrib.gis.geos.point import Point

from datasets.hr.models import BetrokkenPersonen
from datasets.hr.models import CBS_sbi_section
from datasets.hr.models import DataSelectie
from datasets.hr.models import GeoVestigingen

log = logging.getLogger(__name__)

VESTIGING_FIELDS = (
    'vestigingsnummer', 'naam', 'hoofdvestiging',
    'locatie_type', 'geometrie')

BETROKKENEN_FIELDS = (
    'mac_naam', 'rol', 'naam', 'rechtsvorm', 'functietitel',
    'soortbevoegdheid', 'bevoegde_naam')

MAATSCHAPPELIJKE_ACT_FIELDS = (
    'kvk_nummer', 'datum_aanvang', 'datum_einde')

LOCATIE_BEZOEKADRES_FIELDS = (
    'volledig_adres', 'afgeschermd')
LOCATIE_POSTADRES_FIELDS = (
    'straat_huisnummer', 'huisletter', 'huisnummer', 'huisnummertoevoeging',
    'postcode', 'straatnaam', 'postbus_nummer', 'toevoegingadres',
    'volledig_adres', 'afgeschermd')

PROGRESSREPORT = 10  # report every xx seconds

def to_dict(data: object, fields: tuple, field_prefix=None) -> dict:
    result_dict = {}
    for f in fields:
        if field_prefix:
            f_to = field_prefix + '_' + f
        else:
            f_to = f
        if data:
            value = getattr(data, f)
            if isinstance(value, Point):
                value = [value.x, value.y]
            elif isinstance(value, Decimal):
                value = str(value)
            elif isinstance(value, datetime) or isinstance(value, date):
                value = value.isoformat()
        else:
            value = ''
        result_dict[f_to] = value

    return result_dict

def build_joined_ds_table():
    """
    De dataselectie wordt in 2 stappen geschreven. Eerst wordt de api data
    opgebouwd en in een json met een key per vestiging_id weggeschreven.

    Vervolgens wordt in dataselectie deze file gelezen en worden de
    search keys opgenomen in elastic, waardoor die doorzoekbaar wordt
    in dataselectie.
    """

    with db.connection.cursor() as cursor:
        cursor.execute("TRUNCATE TABLE hr_dataselectie")

    sbi_values = build_sbi_filterkeys()

    betrokken_per_vestiging = (
        BetrokkenPersonen.objects.order_by('vestigingsnummer').iterator())

    # Indien geen betrokkenen -> resultaat is None
    try:
        betrokken = next(betrokken_per_vestiging)
    except StopIteration:
        betrokken = None

    totalrowcount = GeoVestigingen.objects.filter().count()

    log.info('START opbouw dataselectie api als json:  %s', totalrowcount)

    # groupby vestigings nummer
    by_vestiging = groupby(
        GeoVestigingen.objects.filter(bag_vbid__isnull=False).order_by(
            'vestigingsnummer'),
        lambda row: row.vestigingsnummer
    )

    count = 1
    lastreport = time.time()

    for vestigingsnummer, vest_data in by_vestiging:
        count, lastreport = measure_progress(totalrowcount, count, lastreport)

        # verzamel gegevens per vestiging
        vst_sbi, vestiging_dict, naam, bag_vbid, bag_numid, countextra = per_vestiging(
            vestigingsnummer, vest_data, sbi_values)
        count += countextra

        # save json in ds tabel
        betrokken = write_hr_dataselectie(
            vst_sbi, betrokken, vestiging_dict,
            vestigingsnummer, betrokken_per_vestiging,
            naam, bag_vbid, bag_numid)

    log.info('Opbouw dataselectie api als json VOLTOOID')

def per_vestiging(vestigingsnummer, vest_data, sbi_values):
    """
    Read data per vestiging and build dict (to be json) per
    vestiging
    """

    first = True
    vst_sbi = []
    vestiging_dict = {}
    sbi_repeat = None
    sbi_double_check = []
    count = 1

    for sbi_repeat in vest_data:

        if first:
            vestiging_dict = to_dict(sbi_repeat, VESTIGING_FIELDS)
            vestiging_dict = add_adressen_dict(vestiging_dict, sbi_repeat)
            vestiging_dict['sbi_codes'] = vst_sbi = []
            first = False

        if not sbi_repeat.sbi_code in sbi_double_check:
            sbi = sbi_values[sbi_repeat.sbi_code]
            sbi['vestigingsnummer'] = vestigingsnummer
            sbi['bedrijfsnaam'] = sbi_repeat.naam
            sbi_double_check.append(sbi_repeat.sbi_code)

            vst_sbi.append(sbi)
        count += 1

    if sbi_repeat.bezoekadres.bag_vbid:
        bag_vbid = sbi_repeat.bezoekadres.bag_vbid
    else:
        bag_vbid = sbi_repeat.postadres.bag_vbid
    return (
        vst_sbi, vestiging_dict, sbi_repeat.naam,
        bag_vbid, sbi_repeat.bezoekadres.bag_numid,
        count)

def write_hr_dataselectie(
        vst_sbi, betrokken, vestiging_dict, vestigingsnummer,
        betrokken_per_vestiging, naam, bag_vbid, bag_numid):
    if len(vst_sbi):
        vestiging_dict, betrokken = add_betrokkenen_to_vestigingen(
            betrokken, vestiging_dict,
            vestigingsnummer, betrokken_per_vestiging)
    else:
        log.error('Vestiging %s %s zonder sbi code' % (vestigingsnummer, naam))

    if vestiging_dict and bag_numid:
        ds = DataSelectie(
            vestigingsnummer, bag_vbid, bag_numid, vestiging_dict)
        ds.save()
    return betrokken

def measure_progress(totalrowcount, count, lastreport):
    """
    Every PROGRESSREPORT timedelta print progres line
    """

    if time.time() - lastreport > PROGRESSREPORT:
        lastreport = time.time()
        pct_complete = str(int((count * 100) / totalrowcount)) + '%'

        log.info('Opbouw dataselectie %s voltooid, rownr %s',
                 pct_complete, count)

    return count, lastreport

def add_adressen_dict(vestiging_dict: dict, sbi_repeat) -> dict:
    """
    Voeg bezoek- en postadressen toe als dict aan resulterende json.
    bezoekadres en postadres zijn relaties van geolocatie, de gevonden rijen
    worden vertaald naar dict. De bezoekadres gegevens worden later in
    dataselectie uit Bag opgehaald. Postadres wordt wel volledig meegegeven
    """

    vestiging_dict.update(
        to_dict(
            sbi_repeat.bezoekadres, LOCATIE_BEZOEKADRES_FIELDS, 'bezoekadres'))

    vestiging_dict['bezoekadres_correctie'] = \
        correctie_address(sbi_repeat.bezoekadres)

    vestiging_dict.update(
        to_dict(
            sbi_repeat.postadres,
            LOCATIE_POSTADRES_FIELDS, 'postadres'))

    vestiging_dict['postadres_correctie'] = \
        correctie_address(sbi_repeat.postadres)

    if vestiging_dict['postadres_correctie'] == 'BAG':
        vestiging_dict.update(
            get_bag(
                sbi_repeat.postadres,
                LOCATIE_POSTADRES_FIELDS, 'postadres'))

    return vestiging_dict

def get_bag(postadres, pa_fields, prefix):
    return {}

def correctie_address(address):
    if address and address.correctie:
        return 'BAG'
    elif address.correctie is False:
        return 'ERR'
    else:
        return 'KVK'

def add_betrokkenen_to_vestigingen(
        betrokken, vestiging_dict, vestigingsnummer,
        betrokken_per_vestiging):
    """
    Voeg betrokkenen aan vestiging toe. Let op, dit is een balanceline
    verwerking. Om de snelheid hoog te houden lees ik de betrokkenen
    per vestiging in en verwerk ze voor de vestiging waar ik mee
    bezig ben. Vergelijkbaar met een join op SQL.
    De oplossing is gekozen vanwege complexiteit (en traagheid) van
    de resulterende join in SQL tegenover de snelheid van sequentiele
    verwerking
    """

    vestiging_dict['betrokkenen'] = vst_betr = []
    for field in MAATSCHAPPELIJKE_ACT_FIELDS:
        vestiging_dict[field] = None
    first = True

    while betrokken and betrokken.vestigingsnummer <= vestigingsnummer:
        if betrokken.vestigingsnummer == vestigingsnummer:
            vst_betr.append(to_dict(betrokken, BETROKKENEN_FIELDS))
            if first:
                mac_dict = to_dict(betrokken, MAATSCHAPPELIJKE_ACT_FIELDS)
                vestiging_dict.update(mac_dict)
                first = False

        if betrokken.vestigingsnummer <= vestigingsnummer:
            try:
                betrokken = next(betrokken_per_vestiging)
            except StopIteration:
                break

    return vestiging_dict, betrokken

def build_sbi_filterkeys():
    """
    One time build of sbi codes to enable rapid check against
    parameters (aggs)
    :return:
    """
    sbi_flat_values = {}
    for hoofdcat in CBS_sbi_section.objects.select_related():
        for subcat in hoofdcat.cbs_sbi_rootnode_set.all():
            for sbi in subcat.cbs_sbicode_set.all():
                sbi_flat_values[sbi.sbi_code] = \
                    {
                        'sbi_code': sbi.sbi_code,
                        'hoofdcategorie': hoofdcat.title,
                        'sub_sub_categorie': sbi.title,
                        'subcategorie': subcat.title,
                        'hcat': hoofdcat.code,
                        'scat': subcat.code
                    }
    return sbi_flat_values
<commit_msg>Voorkom foute data in dataselectie<commit_after>import logging
import time
from datetime import datetime, date
from decimal import Decimal
from itertools import groupby

from django import db
from django.contrib.gis.geos.point import Point

from datasets.hr.models import BetrokkenPersonen
from datasets.hr.models import CBS_sbi_section
from datasets.hr.models import DataSelectie
from datasets.hr.models import GeoVestigingen

log = logging.getLogger(__name__)

VESTIGING_FIELDS = (
    'vestigingsnummer', 'naam', 'hoofdvestiging',
    'locatie_type', 'geometrie')

BETROKKENEN_FIELDS = (
    'mac_naam', 'rol', 'naam', 'rechtsvorm', 'functietitel',
    'soortbevoegdheid', 'bevoegde_naam')

MAATSCHAPPELIJKE_ACT_FIELDS = (
    'kvk_nummer', 'datum_aanvang', 'datum_einde')

LOCATIE_BEZOEKADRES_FIELDS = (
    'volledig_adres', 'afgeschermd')
LOCATIE_POSTADRES_FIELDS = (
    'straat_huisnummer', 'huisletter', 'huisnummer', 'huisnummertoevoeging',
    'postcode', 'straatnaam', 'postbus_nummer', 'toevoegingadres',
    'volledig_adres', 'afgeschermd')

PROGRESSREPORT = 10  # report every xx seconds

def to_dict(data: object, fields: tuple, field_prefix=None) -> dict:
    result_dict = {}
    for f in fields:
        if field_prefix:
            f_to = field_prefix + '_' + f
        else:
            f_to = f
        if data:
            value = getattr(data, f)
            if isinstance(value, Point):
                value = [value.x, value.y]
            elif isinstance(value, Decimal):
                value = str(value)
            elif isinstance(value, datetime) or isinstance(value, date):
                value = value.isoformat()
        else:
            value = ''
        result_dict[f_to] = value

    return result_dict

def build_joined_ds_table():
    """
    De dataselectie wordt in 2 stappen geschreven. Eerst wordt de api data
    opgebouwd en in een json met een key per vestiging_id weggeschreven.

    Vervolgens wordt in dataselectie deze file gelezen en worden de
    search keys opgenomen in elastic, waardoor die doorzoekbaar wordt
    in dataselectie.
    """

    with db.connection.cursor() as cursor:
        cursor.execute("TRUNCATE TABLE hr_dataselectie")

    sbi_values = build_sbi_filterkeys()

    betrokken_per_vestiging = (
        BetrokkenPersonen.objects.order_by('vestigingsnummer').iterator())

    # Indien geen betrokkenen -> resultaat is None
    try:
        betrokken = next(betrokken_per_vestiging)
    except StopIteration:
        betrokken = None

    totalrowcount = GeoVestigingen.objects.filter().count()

    log.info('START opbouw dataselectie api als json:  %s', totalrowcount)

    # groupby vestigings nummer
    by_vestiging = groupby(
        GeoVestigingen.objects.filter(bag_vbid__isnull=False).order_by(
            'vestigingsnummer'),
        lambda row: row.vestigingsnummer
    )

    count = 1
    lastreport = time.time()

    for vestigingsnummer, vest_data in by_vestiging:
        count, lastreport = measure_progress(totalrowcount, count, lastreport)

        # verzamel gegevens per vestiging
        vst_sbi, vestiging_dict, naam, bag_vbid, bag_numid, countextra = per_vestiging(
            vestigingsnummer, vest_data, sbi_values)
        count += countextra

        # save json in ds tabel
        betrokken = write_hr_dataselectie(
            vst_sbi, betrokken, vestiging_dict,
            vestigingsnummer, betrokken_per_vestiging,
            naam, bag_vbid, bag_numid)

    log.info('Opbouw dataselectie api als json VOLTOOID')

def per_vestiging(vestigingsnummer, vest_data, sbi_values):
    """
    Read data per vestiging and build dict (to be json) per
    vestiging
    """

    first = True
    vst_sbi = []
    vestiging_dict = {}
    sbi_repeat = None
    sbi_double_check = []
    count = 1

    for sbi_repeat in vest_data:

        if first:
            vestiging_dict = to_dict(sbi_repeat, VESTIGING_FIELDS)
            vestiging_dict = add_adressen_dict(vestiging_dict, sbi_repeat)
            vestiging_dict['sbi_codes'] = vst_sbi = []
            first = False

        if not sbi_repeat.sbi_code in sbi_double_check:
            sbi = sbi_values[sbi_repeat.sbi_code]
            sbi['vestigingsnummer'] = vestigingsnummer
            sbi['bedrijfsnaam'] = sbi_repeat.naam
            sbi_double_check.append(sbi_repeat.sbi_code)

            vst_sbi.append(sbi)
        count += 1

    if sbi_repeat.bezoekadres.bag_vbid:
        bag_vbid = sbi_repeat.bezoekadres.bag_vbid
    else:
        bag_vbid = sbi_repeat.postadres.bag_vbid
    return (
        vst_sbi, vestiging_dict, sbi_repeat.naam,
        bag_vbid, sbi_repeat.bezoekadres.bag_numid,
        count)

def write_hr_dataselectie(
        vst_sbi, betrokken, vestiging_dict, vestigingsnummer,
        betrokken_per_vestiging, naam, bag_vbid, bag_numid):
    if len(vst_sbi):
        vestiging_dict, betrokken = add_betrokkenen_to_vestigingen(
            betrokken, vestiging_dict,
            vestigingsnummer, betrokken_per_vestiging)
    else:
        log.error('Vestiging %s %s zonder sbi code' % (vestigingsnummer, naam))

    if vestiging_dict and bag_vbid:
        ds = DataSelectie(
            vestigingsnummer, bag_vbid, bag_numid, vestiging_dict)
        ds.save()
    return betrokken

def measure_progress(totalrowcount, count, lastreport):
    """
    Every PROGRESSREPORT timedelta print progres line
    """

    if time.time() - lastreport > PROGRESSREPORT:
        lastreport = time.time()
        pct_complete = str(int((count * 100) / totalrowcount)) + '%'

        log.info('Opbouw dataselectie %s voltooid, rownr %s',
                 pct_complete, count)

    return count, lastreport

def add_adressen_dict(vestiging_dict: dict, sbi_repeat) -> dict:
    """
    Voeg bezoek- en postadressen toe als dict aan resulterende json.
    bezoekadres en postadres zijn relaties van geolocatie, de gevonden rijen
    worden vertaald naar dict. De bezoekadres gegevens worden later in
    dataselectie uit Bag opgehaald. Postadres wordt wel volledig meegegeven
    """

    vestiging_dict.update(
        to_dict(
            sbi_repeat.bezoekadres, LOCATIE_BEZOEKADRES_FIELDS, 'bezoekadres'))

    vestiging_dict['bezoekadres_correctie'] = \
        correctie_address(sbi_repeat.bezoekadres)

    vestiging_dict.update(
        to_dict(
            sbi_repeat.postadres,
            LOCATIE_POSTADRES_FIELDS, 'postadres'))

    vestiging_dict['postadres_correctie'] = \
        correctie_address(sbi_repeat.postadres)

    if vestiging_dict['postadres_correctie'] == 'BAG':
        vestiging_dict.update(
            get_bag(
                sbi_repeat.postadres,
                LOCATIE_POSTADRES_FIELDS, 'postadres'))

    return vestiging_dict

def get_bag(postadres, pa_fields, prefix):
    return {}

def correctie_address(address):
    if address and address.correctie:
        return 'BAG'
    elif address.correctie is False:
        return 'ERR'
    else:
        return 'KVK'

def add_betrokkenen_to_vestigingen(
        betrokken, vestiging_dict, vestigingsnummer,
        betrokken_per_vestiging):
    """
    Voeg betrokkenen aan vestiging toe. Let op, dit is een balanceline
    verwerking. Om de snelheid hoog te houden lees ik de betrokkenen
    per vestiging in en verwerk ze voor de vestiging waar ik mee
    bezig ben. Vergelijkbaar met een join op SQL.
    De oplossing is gekozen vanwege complexiteit (en traagheid) van
    de resulterende join in SQL tegenover de snelheid van sequentiele
    verwerking
    """

    vestiging_dict['betrokkenen'] = vst_betr = []
    for field in MAATSCHAPPELIJKE_ACT_FIELDS:
        vestiging_dict[field] = None
    first = True

    while betrokken and betrokken.vestigingsnummer <= vestigingsnummer:
        if betrokken.vestigingsnummer == vestigingsnummer:
            vst_betr.append(to_dict(betrokken, BETROKKENEN_FIELDS))
            if first:
                mac_dict = to_dict(betrokken, MAATSCHAPPELIJKE_ACT_FIELDS)
                vestiging_dict.update(mac_dict)
                first = False

        if betrokken.vestigingsnummer <= vestigingsnummer:
            try:
                betrokken = next(betrokken_per_vestiging)
            except StopIteration:
                break

    return vestiging_dict, betrokken

def build_sbi_filterkeys():
    """
    One time build of sbi codes to enable rapid check against
    parameters (aggs)
    :return:
    """
    sbi_flat_values = {}
    for hoofdcat in CBS_sbi_section.objects.select_related():
        for subcat in hoofdcat.cbs_sbi_rootnode_set.all():
            for sbi in subcat.cbs_sbicode_set.all():
                sbi_flat_values[sbi.sbi_code] = \
                    {
                        'sbi_code': sbi.sbi_code,
                        'hoofdcategorie': hoofdcat.title,
                        'sub_sub_categorie': sbi.title,
                        'subcategorie': subcat.title,
                        'hcat': hoofdcat.code,
                        'scat': subcat.code
                    }
    return sbi_flat_values
<|endoftext|>


<commit_before>import cv2

WINDOW_NAME = 'hello'

def global_on_mouse(event, x, y, unknown, lego_player):
    lego_player.on_mouse(event, x, y)

class LegoPlayer(object):
    def __init__(self):
        self.rect = []

        cv2.namedWindow(WINDOW_NAME)
        cv2.setMouseCallback(WINDOW_NAME, global_on_mouse, self)
        self.capture = cv2.VideoCapture(0)

    def on_mouse(self, event, x, y):
        if event == cv2.EVENT_LBUTTONDOWN:
            if len(self.rect) == 4:
                self.rect = []
            self.rect.append((x, y))
            print self.rect

    def process_frame(self, frame):
        return cv2.split(frame)[2]

    def loop(self):
        while True:
            success, frame = self.capture.read()
            result = self.process_frame(frame)
            cv2.imshow(WINDOW_NAME, result)
            cv2.waitKey(10)

if __name__ == '__main__':
    lego_player = LegoPlayer()
    lego_player.loop()
<commit_msg>Use numpy array for ROI<commit_after>import numpy as np
import cv2

WINDOW_NAME = 'hello'

def global_on_mouse(event, x, y, unknown, lego_player):
    lego_player.on_mouse(event, x, y)

class LegoPlayer(object):
    def __init__(self):
        self.rect = np.empty((4, 2))
        self.rect_index = -1

        cv2.namedWindow(WINDOW_NAME)
        cv2.setMouseCallback(WINDOW_NAME, global_on_mouse, self)
        self.capture = cv2.VideoCapture(0)

    def has_roi(self):
        return self.rect_index == 3

    def on_mouse(self, event, x, y):
        if event == cv2.EVENT_LBUTTONDOWN:
            if self.has_roi():
                self.rect_index = -1
            self.rect_index += 1
            self.rect[self.rect_index][0] = x
            self.rect[self.rect_index][1] = y
            if self.has_roi():
                print self.rect

    def process_frame(self, frame):
        return cv2.split(frame)[2]

    def loop(self):
        while True:
            success, frame = self.capture.read()
            result = self.process_frame(frame)
            cv2.imshow(WINDOW_NAME, result)
            cv2.waitKey(10)

if __name__ == '__main__':
    lego_player = LegoPlayer()
    lego_player.loop()
<|endoftext|>


<commit_before>import argparse
import os
import re
from markdown import Markdown
import importlib
try:
    import controller
except ImportError:
    controller = None

##############
# Public API #
##############

def link(text):
    """Link the text based on <include> tags.

    An include tag should have the following syntax:

        <include path/to/file.ext>

    The file can be of any filetype, and include tags can have paths
    to any filetype (even filetypes that are different from the
    original source file). Using the above syntax, all of the contents
    of the included file will be substituted for the tag.

    There is an alternate syntax:

        <include path/to/file.ext:blockName>

    The block specified by blockName should be defined using the
    <block> tag in file.ext. Only the contents within blockName will
    be included.

    Note that blockName is a specific use case of a general pattern:

        <include path/to/file.ext:blockRegex>

    where blockRegex is a regular expression that matches block
    names. Each block that is matched by blockRegex acts as if it
    had been written in its own include tag. For example, if the
    regular expression "block\d" matches three blocks "block1",
    "block2", and "block3", then the include tag will be expanded into

        <include path/to/file.ext:block1>
        <include path/to/file.ext:block2>
        <include path/to/file.ext:block3>

    Note: the regular expression must match the entire block name. For
    example, if blockRegex was "block\d", it would not match a block
    called "block10".

    RETURNS:
    text -- str; the text after resolving all include tags. Any block
            tags that are left over are preserved in text.
    """
    return substitute_links(text, {})

def retrieve_blocks(text):
    """Strip block tags from text and return a mapping of block
    names to block contents.

    RETURNS:
    text     -- str; the text with all block tags removed
    mappings -- dict; mapping of block names to block contents
                (strings)
    """
    cache = {}
    return cache_blocks('', text, cache), cache

def scrape_headers(text, regex, translate):
    """Scrape headers based on a regular expression and substitution.

    PARAMETERS:
    text      -- str; the source text
    regex     -- str or RegexObject; a regular expression that defines
                 the semantics of a header
    translate -- function; a function that takes in a Match object.
                 The translation function converts the match into
                 a final representation of the header (i.e. the
                 return value will typically be a string, but can
                 be any object).

    regex and translate provide some flexibility in defining what a
    "header" is. For example, a header can be defined as all
    <h[0-6]> tags in an HTML document to create a table of contents.
    Alternatively, a header can be defined as a function signature
    in a Python script.

    RETURNS:
    headers -- list; headers, as converted by the translate function.
               the contents of this list can be any data type; they
               will appear in the order that the original headers
               are placed within text
    """
    # TODO

def substitutions(text, subs):
    """Apply regular expression substitutions defined for custom
    patterns.

    PARAMETERS:
    text -- str; original source text
    subs -- list of 2-tuples (regex, sub); regex is either a str or a
            RegexObject, which defines the pattern. sub is a
            substitution function, as used by re.sub

    The (regex, sub) pairs will be applied in the order that they
    appear in subs. This means it is the responsibility of the caller
    of this function to select a proper ordering of substitutions.

    RETURNS:
    text -- str; the text after applying all substitutions in subs
    """
    # TODO

def postprocess(text):
    cache = {}
    converted = Markdown(text, pre_hook=link)
    for k, v in converted.variables.items():
        cache[k] = v
    if controller:
        text = apply_controller(converted.text)
        for k, v in controller.configs.items():
            cache[k] = v
        toc = scrape_toc(converted.text)
        cache['table-of-contents'] = controller.toc(toc)
    return cache_blocks('', converted.text, cache), cache

# Markdown(text, pre_hook=link)

def link_markdown(text):
    cache = {}
    converted = Markdown(text, pre_hook=link)
    for k, v in converted.variables.items():
        cache[k] = v
    if controller:
        text = apply_controller(converted.text)
        for k, v in controller.configs.items():
            cache[k] = v
        toc = scrape_toc(converted.text)
        cache['table-of-contents'] = controller.toc(toc)
    return cache_blocks('', converted.text, cache), cache

##########
# Linker #
##########

re_include = re.compile(r"""
    <\s*
        include
        \s+
        (.+?)       # \1 is filename
        (?:
            :       # colon as delimiter
            (.+?)   # \2 is block regex
        )?          # optional
    \s*>
""", re.X)
def substitute_links(text, cache):
    """Create a function that acts as a substitution function for the
    include-tag regex.
    """
    def link_sub(match):
        filename = match.group(1)
        regex = match.group(2)
        text = retrieve_and_link(filename, cache)
        if not regex:
            return cache[filename + ':all']
        return '\n'.join([cache[block] for block in cache
                          if re.match('^' + regex + '$',
                                      block.split(':')[-1])])
    return re_include.sub(link_sub, text)

def retrieve_and_link(filename, cache):
    with open(filename, 'r') as f:
        text = f.read()
    text = substitute_links(text, cache)
    cache_blocks(filename, text, cache)
    return text

re_block = re.compile(r"""
    [^\n]*
    <\s*
        block
        \s+
        (.+?)       # \1 is block name
    \s*>
    [^\n]*\n
    (.*?)           # \2 is block contents
    \n[^\n]*
    </\s*        # forward slash to denote closing tag
        block
        \s+
        \1
    \s*>
    [^\n]*
""", re.S | re.X)

def cache_blocks(filename, text, cache):
    while re_block.search(text):
        for name, contents in re_block.findall(text):
            contents = cache_blocks(filename, contents, cache)
            cache[filename + ':' + name] = contents
        text = re_block.sub(lambda m: m.group(2), text)
    cache[filename + ':all'] = text
    return text

def apply_controller(text):
    for regex, sub in controller.regexes:
        text = regex.sub(sub, text)
    return text

header_re = re.compile(r"""
<\s*h([1-6])(?:.*?id=(['"])(.*?)\2.*?)?>
(.*?)
<\s*/\s*h\1\s*>
""", re.X)
def scrape_toc(text):
    return [(h[0], h[2], h[3]) for h in header_re.findall(text)]

##########################
# Command-line Interface #
##########################

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str,
                        help="Link contents of file")
    parser.add_argument('-d', '--destination', type=str,
                        help="Store result in destination file")
    parser.add_argument('-c', '--cache', action='store_true',
                        help="Show cache keys")
    args = parser.parse_args()

    if not os.path.exists(args.file):
        print('File ' + args.file + ' does not exist.')
        exit(1)
    elif not os.path.isfile(args.file):
        print(args.file + ' is not a valid file')
        exit(1)
    with open(args.file, 'r') as f:
        result = link(f.read())
    if args.cache:
        result, cache = retrieve_blocks(result)
        print('--- Cache keys ---')
        for k in sorted(cache):
            print(k)
    if args.destination:
        with open(args.destination, 'w') as f:
            f.write(result)
        print('Result can be found in ' + args.destination)
    else:
        print('--- BEGIN RESULT ---')
        print(result)
        print('--- END RESULT ---')

if __name__ == '__main__':
    main()

<commit_msg>Implement substitutions and scrape_headers<commit_after>import argparse
import os
import re
from markdown import Markdown
import importlib
try:
    import controller
except ImportError:
    controller = None

##############
# Public API #
##############

def link(text):
    """Link the text based on <include> tags.

    An include tag should have the following syntax:

        <include path/to/file.ext>

    The file can be of any filetype, and include tags can have paths
    to any filetype (even filetypes that are different from the
    original source file). Using the above syntax, all of the contents
    of the included file will be substituted for the tag.

    There is an alternate syntax:

        <include path/to/file.ext:blockName>

    The block specified by blockName should be defined using the
    <block> tag in file.ext. Only the contents within blockName will
    be included.

    Note that blockName is a specific use case of a general pattern:

        <include path/to/file.ext:blockRegex>

    where blockRegex is a regular expression that matches block
    names. Each block that is matched by blockRegex acts as if it
    had been written in its own include tag. For example, if the
    regular expression "block\d" matches three blocks "block1",
    "block2", and "block3", then the include tag will be expanded into

        <include path/to/file.ext:block1>
        <include path/to/file.ext:block2>
        <include path/to/file.ext:block3>

    Note: the regular expression must match the entire block name. For
    example, if blockRegex was "block\d", it would not match a block
    called "block10".

    RETURNS:
    text -- str; the text after resolving all include tags. Any block
            tags that are left over are preserved in text.
    """
    return substitute_links(text, {})

def retrieve_blocks(text):
    """Strip block tags from text and return a mapping of block
    names to block contents.

    RETURNS:
    text     -- str; the text with all block tags removed
    mappings -- dict; mapping of block names to block contents
                (strings)
    """
    cache = {}
    return cache_blocks('', text, cache), cache

def scrape_headers(text, regex, translate):
    """Scrape headers based on a regular expression and substitution.

    PARAMETERS:
    text      -- str; the source text
    regex     -- str or RegexObject; a regular expression that defines
                 the semantics of a header
    translate -- function; a function that takes in a Match object.
                 The translation function converts the match into
                 a final representation of the header (i.e. the
                 return value will typically be a string, but can
                 be any object).

    regex and translate provide some flexibility in defining what a
    "header" is. For example, a header can be defined as all
    <h[0-6]> tags in an HTML document to create a table of contents.
    Alternatively, a header can be defined as a function signature
    in a Python script.

    RETURNS:
    headers -- list; headers, as converted by the translate function.
               the contents of this list can be any data type; they
               will appear in the order that the original headers
               are placed within text
    """
    return [translate(match) for match in re.findall(regex, text)]

def substitutions(text, subs):
    """Apply regular expression substitutions defined for custom
    patterns.

    PARAMETERS:
    text -- str; original source text
    subs -- list of 2-tuples (regex, sub); regex is either a str or a
            RegexObject, which defines the pattern. sub is a
            substitution function, as used by re.sub

    The (regex, sub) pairs will be applied in the order that they
    appear in subs. This means it is the responsibility of the caller
    of this function to select a proper ordering of substitutions.

    RETURNS:
    text -- str; the text after applying all substitutions in subs
    """
    for regex, sub in subs:
        text = re.sub(regex, sub, text)
    return text

def postprocess(text):
    cache = {}
    converted = Markdown(text, pre_hook=link)
    for k, v in converted.variables.items():
        cache[k] = v
    if controller:
        text = apply_controller(converted.text)
        for k, v in controller.configs.items():
            cache[k] = v
        toc = scrape_toc(converted.text)
        cache['table-of-contents'] = controller.toc(toc)
    return cache_blocks('', converted.text, cache), cache

# Markdown(text, pre_hook=link)

def link_markdown(text):
    cache = {}
    converted = Markdown(text, pre_hook=link)
    for k, v in converted.variables.items():
        cache[k] = v
    if controller:
        text = apply_controller(converted.text)
        for k, v in controller.configs.items():
            cache[k] = v
        toc = scrape_toc(converted.text)
        cache['table-of-contents'] = controller.toc(toc)
    return cache_blocks('', converted.text, cache), cache

##########
# Linker #
##########

re_include = re.compile(r"""
    <\s*
        include
        \s+
        (.+?)       # \1 is filename
        (?:
            :       # colon as delimiter
            (.+?)   # \2 is block regex
        )?          # optional
    \s*>
""", re.X)
def substitute_links(text, cache):
    """Create a function that acts as a substitution function for the
    include-tag regex.
    """
    def link_sub(match):
        filename = match.group(1)
        regex = match.group(2)
        text = retrieve_and_link(filename, cache)
        if not regex:
            return cache[filename + ':all']
        return '\n'.join([cache[block] for block in cache
                          if re.match('^' + regex + '$',
                                      block.split(':')[-1])])
    return re_include.sub(link_sub, text)

def retrieve_and_link(filename, cache):
    with open(filename, 'r') as f:
        text = f.read()
    text = substitute_links(text, cache)
    cache_blocks(filename, text, cache)
    return text

re_block = re.compile(r"""
    [^\n]*
    <\s*
        block
        \s+
        (.+?)       # \1 is block name
    \s*>
    [^\n]*\n
    (.*?)           # \2 is block contents
    \n[^\n]*
    </\s*        # forward slash to denote closing tag
        block
        \s+
        \1
    \s*>
    [^\n]*
""", re.S | re.X)

def cache_blocks(filename, text, cache):
    while re_block.search(text):
        for name, contents in re_block.findall(text):
            contents = cache_blocks(filename, contents, cache)
            cache[filename + ':' + name] = contents
        text = re_block.sub(lambda m: m.group(2), text)
    cache[filename + ':all'] = text
    return text

def apply_controller(text):
    for regex, sub in controller.regexes:
        text = regex.sub(sub, text)
    return text

header_re = re.compile(r"""
<\s*h([1-6])(?:.*?id=(['"])(.*?)\2.*?)?>
(.*?)
<\s*/\s*h\1\s*>
""", re.X)
def scrape_toc(text):
    return [(h[0], h[2], h[3]) for h in header_re.findall(text)]

##########################
# Command-line Interface #
##########################

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str,
                        help="Link contents of file")
    parser.add_argument('-d', '--destination', type=str,
                        help="Store result in destination file")
    parser.add_argument('-c', '--cache', action='store_true',
                        help="Show cache keys")
    args = parser.parse_args()

    if not os.path.exists(args.file):
        print('File ' + args.file + ' does not exist.')
        exit(1)
    elif not os.path.isfile(args.file):
        print(args.file + ' is not a valid file')
        exit(1)
    with open(args.file, 'r') as f:
        result = link(f.read())
    if args.cache:
        result, cache = retrieve_blocks(result)
        print('--- Cache keys ---')
        for k in sorted(cache):
            print(k)
    if args.destination:
        with open(args.destination, 'w') as f:
            f.write(result)
        print('Result can be found in ' + args.destination)
    else:
        print('--- BEGIN RESULT ---')
        print(result)
        print('--- END RESULT ---')

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before># encoding: utf-8
"""
protocol.py

Created by Thomas Mangin on 2009-08-25.
Copyright (c) 2009-2013 Exa Networks. All rights reserved.
"""

import os

from exabgp.rib.table import Table
from exabgp.rib.delta import Delta

from exabgp.reactor.network.outgoing import Outgoing
from exabgp.reactor.network.error import SizeError,NotifyError

from exabgp.bgp.message import Message
from exabgp.bgp.message.nop import NOP
from exabgp.bgp.message.unknown import Unknown
from exabgp.bgp.message.open import Open
from exabgp.bgp.message.open.capability import Capabilities
from exabgp.bgp.message.open.capability.negotiated import Negotiated
from exabgp.bgp.message.update import Update
from exabgp.bgp.message.update.eor import EOR
from exabgp.bgp.message.keepalive import KeepAlive
from exabgp.bgp.message.notification import Notification, Notify
from exabgp.bgp.message.refresh import RouteRefresh

from exabgp.reactor.api.processes import ProcessError

from exabgp.logger import Logger,FakeLogger,LazyFormat

# This is the number of chuncked message we are willing to buffer, not the number of routes
MAX_BACKLOG = 15000

_NOP = NOP()
_UPDATE = Update()
_UNKNOWN = Unknown()

class Protocol (object):
	decode = True

	def __init__ (self,peer):
		try:
			self.logger = Logger()
		except RuntimeError:
			self.logger = FakeLogger()
		self.peer = peer
		self.neighbor = peer.neighbor
		self.negotiated = Negotiated()
		self.delta = Delta(Table(peer))
		self.connection = None
		port = os.environ.get('exabgp.tcp.port','')
		self.port = int(port) if port.isdigit() else 179

		# XXX: FIXME: check the the -19 is correct (but it is harmless)
		# The message size is the whole BGP message _without_ headers
		self.message_size = Message.MAX_LEN-Message.HEADER_LEN

	# XXX: we use self.peer.neighbor.peer_address when we could use self.neighbor.peer_address

	def me (self,message):
		return "Peer %15s ASN %-7s %s" % (self.peer.neighbor.peer_address,self.peer.neighbor.peer_as,message)

	def accept (self,incoming):
		if not self.connection:
			self.connection = incoming

		if self.peer.neighbor.api.neighbor_changes:
			self.peer.reactor.processes.connected(self.peer.neighbor.peer_address)

	def connect (self):
		# allows to test the protocol code using modified StringIO with a extra 'pending' function
		if not self.connection:
			peer = self.neighbor.peer_address
			local = self.neighbor.local_address
			md5 = self.neighbor.md5
			ttl = self.neighbor.ttl
			self.connection = Outgoing(peer.afi,peer.ip,local.ip,self.port,md5,ttl)

			if self.peer.neighbor.api.neighbor_changes:
				self.peer.reactor.processes.connected(self.peer.neighbor.peer_address)

	def close (self,reason='unspecified'):
		self.logger.info(self.me(reason))
		if self.connection:
			# must be first otherwise we could have a loop caused by the raise in the below
			self.connection.close()
			self.connection = None

			try:
				if self.peer.neighbor.api.neighbor_changes:
					self.peer.reactor.processes.down(self.peer.neighbor.peer_address,reason)
			except ProcessError:
				self.logger.message(self.me('could not send notification of neighbor close to API'))

	def write (self,message):
		if self.neighbor.api.send_packets:
			self.peer.reactor.processes.send(self.peer.neighbor.peer_address,message[18],message[:19],message[19:])
		for boolean in self.connection.writer(message):
			yield boolean

	# Read from network .......................................................

	def read_message (self,comment=''):
		try:
			for length,msg,header,body in self.connection.reader():
				if not length:
					yield _NOP
		except NotifyError,n:
			raise Notify(n.code,n.subcode,str(n))

		if self.neighbor.api.receive_packets:
			self.peer.reactor.processes.receive(self.peer.neighbor.peer_address,msg,header,body)

		if msg == Message.Type.UPDATE:
			self.logger.message(self.me('<< UPDATE'))

			if length == 30 and body.startswith(EOR.PREFIX):
				yield EOR().factory(body)

			if self.neighbor.api.receive_routes:
				update = Update().factory(self.negotiated,body)

				for route in update.routes:
					self.logger.routes(LazyFormat(self.me(''),str,route))

				self.peer.reactor.processes.routes(self.neighbor.peer_address,update.routes)
				yield update
			else:
				yield _UNKNOWN

		elif msg == Message.Type.KEEPALIVE:
			self.logger.message(self.me('<< KEEPALIVE%s' % (' (%s)' % comment if comment else '')))
			yield KeepAlive()

		elif msg == Message.Type.NOTIFICATION:
			self.logger.message(self.me('<< NOTIFICATION'))
			yield Notification().factory(body)

		elif msg == Message.Type.ROUTE_REFRESH:
			self.logger.message(self.me('<< ROUTE-REFRESH'))
			yield RouteRefresh().factory(body)

		elif msg == Message.Type.OPEN:
			yield Open().factory(body)

		else:
			self.logger.message(self.me('<< NOP (unknow type %d)' % msg))
			yield Unknown().factory(msg)

	def validate_open (self):
		error = self.negotiated.validate(self.neighbor)
		if error is not None:
			raise Notify(*error)

	def read_open (self,ip):
		for received_open in self.read_message():
			if received_open.TYPE == NOP.TYPE:
				yield received_open
			else:
				break

		if received_open.TYPE != Open.TYPE:
			raise Notify(5,1,'The first packet recevied is not an open message (%s)' % received_open)

		self.logger.message(self.me('<< %s' % received_open))
		yield received_open

	def read_keepalive (self,comment=''):
		for message in self.read_message(comment):
			if message.TYPE == NOP.TYPE:
				yield message
			else:
				break

		if message.TYPE != KeepAlive.TYPE:
			raise Notify(5,2)

		yield message

	#
	# Sending message to peer
	#

	def new_open (self,restarted):
		sent_open = Open().new(
			4,
			self.neighbor.local_as,
			self.neighbor.router_id.ip,
			Capabilities().new(self.neighbor,restarted),
			self.neighbor.hold_time
		)

		# we do not buffer open message in purpose
		for _ in self.write(sent_open.message()):
			yield _NOP

		self.logger.message(self.me('>> %s' % sent_open))
		yield sent_open

	def new_keepalive (self,comment=''):
		keepalive = KeepAlive()

		for _ in self.write(keepalive.message()):
			yield _NOP

		self.logger.message(self.me('>> KEEPALIVE%s' % (' (%s)' % comment if comment else '')))

		yield keepalive

	def new_notification (self,notification):
		for _ in self.write(notification.message()):
			yield _NOP
		self.logger.error(self.me('>> NOTIFICATION (%d,%d,"%s")' % (notification.code,notification.subcode,notification.data)))
		yield notification

	def new_update (self):
		# XXX: This should really be calculated once only
		for _ in self._announce('UPDATE',self.peer.proto.delta.updates(self.negotiated,self.neighbor.group_updates)):
			yield _NOP
		yield _UPDATE

	def new_eors (self):
		eor = EOR().new(self.negotiated.families)
		for _ in self._announce(str(eor),eor.updates(self.negotiated)):
			yield _NOP
		yield _UPDATE

	def _announce (self,name,generator):
		def chunked (generator,size):
			chunk = ''
			number = 0
			for data in generator:
				if len(data) > size:
					raise SizeError('Can not send BGP update larger than %d bytes on this connection.' % size)
				if len(chunk) + len(data) <= size:
					chunk += data
					number += 1
					continue
				yield number,chunk
				chunk = data
				number = 1
			if chunk:
				yield number,chunk

		for number,update in chunked(generator,self.message_size):
			for boolean in self.write(update):
				if boolean:
					self.logger.message(self.me('>> %d %s(s)' % (number,name)))
					yield number
				else:
					yield 0
<commit_msg>fix a bug where a factory was not converted to new API<commit_after># encoding: utf-8
"""
protocol.py

Created by Thomas Mangin on 2009-08-25.
Copyright (c) 2009-2013 Exa Networks. All rights reserved.
"""

import os

from exabgp.rib.table import Table
from exabgp.rib.delta import Delta

from exabgp.reactor.network.outgoing import Outgoing
from exabgp.reactor.network.error import SizeError,NotifyError

from exabgp.bgp.message import Message
from exabgp.bgp.message.nop import NOP
from exabgp.bgp.message.unknown import Unknown
from exabgp.bgp.message.open import Open
from exabgp.bgp.message.open.capability import Capabilities
from exabgp.bgp.message.open.capability.negotiated import Negotiated
from exabgp.bgp.message.update import Update,UpdateFactory
from exabgp.bgp.message.update.eor import EOR
from exabgp.bgp.message.keepalive import KeepAlive
from exabgp.bgp.message.notification import Notification, Notify
from exabgp.bgp.message.refresh import RouteRefresh

from exabgp.reactor.api.processes import ProcessError

from exabgp.logger import Logger,FakeLogger,LazyFormat

# This is the number of chuncked message we are willing to buffer, not the number of routes
MAX_BACKLOG = 15000

_NOP = NOP()
_UPDATE = Update()
_UNKNOWN = Unknown()

class Protocol (object):
	decode = True

	def __init__ (self,peer):
		try:
			self.logger = Logger()
		except RuntimeError:
			self.logger = FakeLogger()
		self.peer = peer
		self.neighbor = peer.neighbor
		self.negotiated = Negotiated()
		self.delta = Delta(Table(peer))
		self.connection = None
		port = os.environ.get('exabgp.tcp.port','')
		self.port = int(port) if port.isdigit() else 179

		# XXX: FIXME: check the the -19 is correct (but it is harmless)
		# The message size is the whole BGP message _without_ headers
		self.message_size = Message.MAX_LEN-Message.HEADER_LEN

	# XXX: we use self.peer.neighbor.peer_address when we could use self.neighbor.peer_address

	def me (self,message):
		return "Peer %15s ASN %-7s %s" % (self.peer.neighbor.peer_address,self.peer.neighbor.peer_as,message)

	def accept (self,incoming):
		if not self.connection:
			self.connection = incoming

		if self.peer.neighbor.api.neighbor_changes:
			self.peer.reactor.processes.connected(self.peer.neighbor.peer_address)

	def connect (self):
		# allows to test the protocol code using modified StringIO with a extra 'pending' function
		if not self.connection:
			peer = self.neighbor.peer_address
			local = self.neighbor.local_address
			md5 = self.neighbor.md5
			ttl = self.neighbor.ttl
			self.connection = Outgoing(peer.afi,peer.ip,local.ip,self.port,md5,ttl)

			if self.peer.neighbor.api.neighbor_changes:
				self.peer.reactor.processes.connected(self.peer.neighbor.peer_address)

	def close (self,reason='unspecified'):
		self.logger.info(self.me(reason))
		if self.connection:
			# must be first otherwise we could have a loop caused by the raise in the below
			self.connection.close()
			self.connection = None

			try:
				if self.peer.neighbor.api.neighbor_changes:
					self.peer.reactor.processes.down(self.peer.neighbor.peer_address,reason)
			except ProcessError:
				self.logger.message(self.me('could not send notification of neighbor close to API'))

	def write (self,message):
		if self.neighbor.api.send_packets:
			self.peer.reactor.processes.send(self.peer.neighbor.peer_address,message[18],message[:19],message[19:])
		for boolean in self.connection.writer(message):
			yield boolean

	# Read from network .......................................................

	def read_message (self,comment=''):
		try:
			for length,msg,header,body in self.connection.reader():
				if not length:
					yield _NOP
		except NotifyError,n:
			raise Notify(n.code,n.subcode,str(n))

		if self.neighbor.api.receive_packets:
			self.peer.reactor.processes.receive(self.peer.neighbor.peer_address,msg,header,body)

		if msg == Message.Type.UPDATE:
			self.logger.message(self.me('<< UPDATE'))

			if length == 30 and body.startswith(EOR.PREFIX):
				yield EOR().factory(body)

			if self.neighbor.api.receive_routes:
				update = UpdateFactory(self.negotiated,body)

				for route in update.routes:
					self.logger.routes(LazyFormat(self.me(''),str,route))

				self.peer.reactor.processes.routes(self.neighbor.peer_address,update.routes)
				yield update
			else:
				yield _UNKNOWN

		elif msg == Message.Type.KEEPALIVE:
			self.logger.message(self.me('<< KEEPALIVE%s' % (' (%s)' % comment if comment else '')))
			yield KeepAlive()

		elif msg == Message.Type.NOTIFICATION:
			self.logger.message(self.me('<< NOTIFICATION'))
			yield Notification().factory(body)

		elif msg == Message.Type.ROUTE_REFRESH:
			self.logger.message(self.me('<< ROUTE-REFRESH'))
			yield RouteRefresh().factory(body)

		elif msg == Message.Type.OPEN:
			yield Open().factory(body)

		else:
			self.logger.message(self.me('<< NOP (unknow type %d)' % msg))
			yield Unknown().factory(msg)

	def validate_open (self):
		error = self.negotiated.validate(self.neighbor)
		if error is not None:
			raise Notify(*error)

	def read_open (self,ip):
		for received_open in self.read_message():
			if received_open.TYPE == NOP.TYPE:
				yield received_open
			else:
				break

		if received_open.TYPE != Open.TYPE:
			raise Notify(5,1,'The first packet recevied is not an open message (%s)' % received_open)

		self.logger.message(self.me('<< %s' % received_open))
		yield received_open

	def read_keepalive (self,comment=''):
		for message in self.read_message(comment):
			if message.TYPE == NOP.TYPE:
				yield message
			else:
				break

		if message.TYPE != KeepAlive.TYPE:
			raise Notify(5,2)

		yield message

	#
	# Sending message to peer
	#

	def new_open (self,restarted):
		sent_open = Open().new(
			4,
			self.neighbor.local_as,
			self.neighbor.router_id.ip,
			Capabilities().new(self.neighbor,restarted),
			self.neighbor.hold_time
		)

		# we do not buffer open message in purpose
		for _ in self.write(sent_open.message()):
			yield _NOP

		self.logger.message(self.me('>> %s' % sent_open))
		yield sent_open

	def new_keepalive (self,comment=''):
		keepalive = KeepAlive()

		for _ in self.write(keepalive.message()):
			yield _NOP

		self.logger.message(self.me('>> KEEPALIVE%s' % (' (%s)' % comment if comment else '')))

		yield keepalive

	def new_notification (self,notification):
		for _ in self.write(notification.message()):
			yield _NOP
		self.logger.error(self.me('>> NOTIFICATION (%d,%d,"%s")' % (notification.code,notification.subcode,notification.data)))
		yield notification

	def new_update (self):
		# XXX: This should really be calculated once only
		for _ in self._announce('UPDATE',self.peer.proto.delta.updates(self.negotiated,self.neighbor.group_updates)):
			yield _NOP
		yield _UPDATE

	def new_eors (self):
		eor = EOR().new(self.negotiated.families)
		for _ in self._announce(str(eor),eor.updates(self.negotiated)):
			yield _NOP
		yield _UPDATE

	def _announce (self,name,generator):
		def chunked (generator,size):
			chunk = ''
			number = 0
			for data in generator:
				if len(data) > size:
					raise SizeError('Can not send BGP update larger than %d bytes on this connection.' % size)
				if len(chunk) + len(data) <= size:
					chunk += data
					number += 1
					continue
				yield number,chunk
				chunk = data
				number = 1
			if chunk:
				yield number,chunk

		for number,update in chunked(generator,self.message_size):
			for boolean in self.write(update):
				if boolean:
					self.logger.message(self.me('>> %d %s(s)' % (number,name)))
					yield number
				else:
					yield 0
<|endoftext|>