<commit_before>from __future__ import unicode_literals

from decimal import Decimal as D

from gittip.testing import Harness
from gittip.models.participant import Participant

class Tests(Harness):

    def make_team(self):
        team = self.make_participant('A Team', number='plural')
        warbucks = self.make_participant('Daddy Warbucks', last_bill_result='')
        warbucks.set_tip_to(team, '100')
        self.warbucks = warbucks
        return team

    def make_participant(self, username, *arg, **kw):
        take_last_week = kw.pop('take_last_week', None)
        participant = Harness.make_participant(self, username, **kw)
        if take_last_week is not None:
            if self.db.one('SELECT * FROM paydays') is None:
                self.db.run("INSERT INTO paydays DEFAULT VALUES")
            self.db.run( "INSERT INTO transfers (timestamp, tipper, tippee, amount) "
                         "VALUES (now(), 'A Team', %(tippee)s, %(amount)s)"
                       , dict(tippee=username, amount=take_last_week,)
                        )
        return participant

    def test_we_can_make_a_team(self):
        team = self.make_team()
        assert team.IS_PLURAL

    def test_random_schmoe_is_not_member_of_team(self):
        team = self.make_team()
        schmoe = self.make_participant('schmoe')
        assert not schmoe.member_of(team)

    def test_team_member_is_team_member(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        assert alice.member_of(team)

    def test_cant_grow_tip_a_lot(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        assert team.set_take_for(alice, D('100.00'), alice) == 80

    def test_take_can_double(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        team.set_take_for(alice, D('80.00'), alice)
        assert team.get_take_for(alice) == 80

    def test_take_can_double_but_not_a_penny_more(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        actual = team.set_take_for(alice, D('80.01'), alice)
        assert actual == 80

    def test_increase_is_based_on_actual_take_last_week(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='20.00')
        team._MixinTeam__set_take_for(alice, D('35.00'), team)
        assert team.set_take_for(alice, D('42.00'), alice) == 40

    def test_if_last_week_is_less_than_a_dollar_can_increase_to_a_dollar(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='0.01', claimed_time='now')
        team.add_member(alice)
        actual = team.set_take_for(alice, D('42.00'), team)
        assert actual == 1

    def test_get_members(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), team)
        members = team.get_members(alice)
        assert len(members) == 2
        assert members[0]['username'] == 'alice'
        assert members[0]['take'] == 42
        assert members[0]['balance'] == 58

    def test_takes_and_receiving_are_updated_correctly(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), alice)
        assert alice.takes == 42
        assert alice.receiving == 42
        self.warbucks.set_tip_to(alice, D('10.00'))
        assert alice.takes == 42
        assert alice.receiving == 52
        team.set_take_for(alice, D('50.00'), alice)
        assert alice.takes == 50
        assert alice.receiving == 60

    def test_changes_to_team_receiving_affect_members_take(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), alice)

        self.warbucks.set_tip_to(team, D('10.00'))  # hard times
        alice = Participant.from_username('alice')
        assert alice.receiving == alice.takes == 10
<commit_msg>Passing test for entanglement of member takes<commit_after>from __future__ import unicode_literals

from decimal import Decimal as D

from gittip.testing import Harness
from gittip.models.participant import Participant

class Tests(Harness):

    def make_team(self):
        team = self.make_participant('A Team', number='plural')
        warbucks = self.make_participant('Daddy Warbucks', last_bill_result='')
        warbucks.set_tip_to(team, '100')
        self.warbucks = warbucks
        return team

    def make_participant(self, username, *arg, **kw):
        take_last_week = kw.pop('take_last_week', None)
        participant = Harness.make_participant(self, username, **kw)
        if take_last_week is not None:
            if self.db.one('SELECT * FROM paydays') is None:
                self.db.run("INSERT INTO paydays DEFAULT VALUES")
            self.db.run( "INSERT INTO transfers (timestamp, tipper, tippee, amount) "
                         "VALUES (now(), 'A Team', %(tippee)s, %(amount)s)"
                       , dict(tippee=username, amount=take_last_week,)
                        )
        return participant

    def test_we_can_make_a_team(self):
        team = self.make_team()
        assert team.IS_PLURAL

    def test_random_schmoe_is_not_member_of_team(self):
        team = self.make_team()
        schmoe = self.make_participant('schmoe')
        assert not schmoe.member_of(team)

    def test_team_member_is_team_member(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        assert alice.member_of(team)

    def test_cant_grow_tip_a_lot(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        assert team.set_take_for(alice, D('100.00'), alice) == 80

    def test_take_can_double(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        team.set_take_for(alice, D('80.00'), alice)
        assert team.get_take_for(alice) == 80

    def test_take_can_double_but_not_a_penny_more(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00')
        team._MixinTeam__set_take_for(alice, D('40.00'), team)
        actual = team.set_take_for(alice, D('80.01'), alice)
        assert actual == 80

    def test_increase_is_based_on_actual_take_last_week(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='20.00')
        team._MixinTeam__set_take_for(alice, D('35.00'), team)
        assert team.set_take_for(alice, D('42.00'), alice) == 40

    def test_if_last_week_is_less_than_a_dollar_can_increase_to_a_dollar(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='0.01', claimed_time='now')
        team.add_member(alice)
        actual = team.set_take_for(alice, D('42.00'), team)
        assert actual == 1

    def test_get_members(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), team)
        members = team.get_members(alice)
        assert len(members) == 2
        assert members[0]['username'] == 'alice'
        assert members[0]['take'] == 42
        assert members[0]['balance'] == 58

    def test_takes_and_receiving_are_updated_correctly(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), alice)
        assert alice.takes == 42
        assert alice.receiving == 42
        self.warbucks.set_tip_to(alice, D('10.00'))
        assert alice.takes == 42
        assert alice.receiving == 52
        team.set_take_for(alice, D('50.00'), alice)
        assert alice.takes == 50
        assert alice.receiving == 60

    def test_changes_to_team_receiving_affect_members_take(self):
        team = self.make_team()
        alice = self.make_participant('alice', take_last_week='40.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), alice)

        self.warbucks.set_tip_to(team, D('10.00'))  # hard times
        alice = Participant.from_username('alice')
        assert alice.receiving == alice.takes == 10

    def test_changes_to_others_take_affects_members_take(self):
        team = self.make_team()

        alice = self.make_participant('alice', take_last_week='30.00', claimed_time='now')
        team.add_member(alice)
        team.set_take_for(alice, D('42.00'), alice)

        bob = self.make_participant('bob', take_last_week='50.00', claimed_time='now')
        team.add_member(bob)
        team.set_take_for(bob, D('60.00'), bob)

        alice = Participant.from_username('alice')
        assert alice.receiving == alice.takes == 40
<|endoftext|>


<commit_before>import copy
from typing import List, Union

from django import forms
from django.conf import settings
from django.forms import (
    BaseForm, BaseInlineFormSet, BaseModelForm, BaseModelFormSet,
)
from django.forms.forms import DeclarativeFieldsMetaclass
from django.forms.models import ModelFormMetaclass
from django.utils.safestring import mark_safe

from .strings import LazyI18nString

class I18nWidget(forms.MultiWidget):
    """
    The default form widget for I18nCharField and I18nTextField. It makes
    use of Django's MultiWidget mechanism and does some magic to save you
    time.
    """
    widget = forms.TextInput

    def __init__(self, locales: List[str], field: forms.Field, attrs=None):
        widgets = []
        self.locales = locales
        self.enabled_locales = locales
        self.field = field
        for lng in self.locales:
            a = copy.copy(attrs) or {}
            a['lang'] = lng
            widgets.append(self.widget(attrs=a))
        super().__init__(widgets, attrs)

    def decompress(self, value) -> List[Union[str, None]]:
        data = []
        first_enabled = None
        any_enabled_filled = False
        if not isinstance(value, LazyI18nString):
            value = LazyI18nString(value)
        for i, lng in enumerate(self.locales):
            dataline = (
                value.data[lng]
                if value is not None and (
                    isinstance(value.data, dict) or isinstance(value.data, LazyI18nString.LazyGettextProxy)
                ) and lng in value.data
                else None
            )
            if lng in self.enabled_locales:
                if not first_enabled:
                    first_enabled = i
                if dataline:
                    any_enabled_filled = True
            data.append(dataline)
        if value and not isinstance(value.data, dict) and not isinstance(value.data, LazyI18nString.LazyGettextProxy):
            data[first_enabled] = value.data
        elif value and not any_enabled_filled:
            data[first_enabled] = value.localize(self.enabled_locales[0])
        return data

    def render(self, name: str, value, attrs=None, renderer=None) -> str:
        if self.is_localized:
            for widget in self.widgets:
                widget.is_localized = self.is_localized
        # value is a list of values, each corresponding to a widget
        # in self.widgets.
        if not isinstance(value, list):
            value = self.decompress(value)
        output = []
        final_attrs = self.build_attrs(attrs or dict())
        id_ = final_attrs.get('id', None)
        for i, widget in enumerate(self.widgets):
            if self.locales[i] not in self.enabled_locales:
                continue
            try:
                widget_value = value[i]
            except IndexError:
                widget_value = None
            if id_:
                final_attrs = dict(
                    final_attrs,
                    id='%s_%s' % (id_, i),
                    title=self.locales[i]
                )
            output.append(widget.render(name + '_%s' % i, widget_value, final_attrs, renderer=renderer))
        return mark_safe(self.format_output(output))

    def format_output(self, rendered_widgets) -> str:
        return '
%s
' % (
            ' i18n-form-single-language' if len(rendered_widgets) <= 1 else '',
            ''.join(rendered_widgets),
        )

class I18nTextInput(I18nWidget):
    """
    The default form widget for I18nCharField. It makes use of Django's MultiWidget
    mechanism and does some magic to save you time.
    """
    widget = forms.TextInput

class I18nTextarea(I18nWidget):
    """
    The default form widget for I18nTextField. It makes use of Django's MultiWidget
    mechanism and does some magic to save you time.
    """
    widget = forms.Textarea

class I18nFormField(forms.MultiValueField):
    """
    The form field that is used by I18nCharField and I18nTextField. It makes use
    of Django's MultiValueField mechanism to create one sub-field per available
    language.

    It contains special treatment to make sure that a field marked as "required" is validated
    as "filled out correctly" if *at least one* translation is filled it. It is never required
    to fill in all of them. This has the drawback that the HTML property ``required`` is set on
    none of the fields as this would lead to irritating behaviour.

    :param locales: An iterable of locale codes that the widget should render a field for. If
                    omitted, fields will be rendered for all languages configured in
                    ``settings.LANGUAGES``.
    :param require_all_fields: A boolean, if set to True field requires all translations to be given.
    """

    def compress(self, data_list) -> LazyI18nString:
        locales = self.locales
        data = {}
        for i, value in enumerate(data_list):
            data[locales[i]] = value
        return LazyI18nString(data)

    def clean(self, value) -> LazyI18nString:
        if isinstance(value, LazyI18nString):
            # This happens e.g. if the field is disabled
            return value
        found = False
        found_all = True
        clean_data = []
        errors = []
        for i, field in enumerate(self.fields):
            try:
                field_value = value[i]
            except (IndexError, TypeError):
                field_value = None
            if field_value not in self.empty_values:
                found = True
            elif field.locale in self.widget.enabled_locales:
                found_all = False
            try:
                clean_data.append(field.clean(field_value))
            except forms.ValidationError as e:
                # Collect all validation errors in a single list, which we'll
                # raise at the end of clean(), rather than raising a single
                # exception for the first error we encounter. Skip duplicates.
                errors.extend(m for m in e.error_list if m not in errors)
        if errors:
            raise forms.ValidationError(errors)
        if self.one_required and not found:
            raise forms.ValidationError(self.error_messages['required'], code='required')
        if self.require_all_fields and not found_all:
            raise forms.ValidationError(self.error_messages['incomplete'], code='incomplete')

        out = self.compress(clean_data)
        self.validate(out)
        self.run_validators(out)
        return out

    def __init__(self, *args, **kwargs):
        fields = []
        defaults = {
            'widget': self.widget,
            'max_length': kwargs.pop('max_length', None),
        }
        self.locales = kwargs.pop('locales', [loc[0] for loc in settings.LANGUAGES])
        self.one_required = kwargs.get('required', True)
        require_all_fields = kwargs.pop('require_all_fields', False)
        kwargs['required'] = False
        kwargs['widget'] = kwargs['widget'](
            locales=self.locales, field=self, **kwargs.pop('widget_kwargs', {})
        )
        defaults.update(**kwargs)
        for lngcode in self.locales:
            defaults['label'] = '%s (%s)' % (defaults.get('label'), lngcode)
            field = forms.CharField(**defaults)
            field.locale = lngcode
            fields.append(field)
        super().__init__(
            fields=fields, require_all_fields=False, *args, **kwargs
        )
        self.require_all_fields = require_all_fields

class I18nFormMixin:
    def __init__(self, *args, **kwargs):
        locales = kwargs.pop('locales', None)
        super().__init__(*args, **kwargs)
        if locales:
            for k, field in self.fields.items():
                if isinstance(field, I18nFormField):
                    field.widget.enabled_locales = locales

class BaseI18nModelForm(I18nFormMixin, BaseModelForm):
    """
    This is a helperclass to construct an I18nModelForm.
    """
    pass

class BaseI18nForm(I18nFormMixin, BaseForm):
    """
    This is a helperclass to construct an I18nForm.
    """
    pass

class I18nForm(BaseI18nForm, metaclass=DeclarativeFieldsMetaclass):
    """
    This is a modified version of Django's Form which differs from Form in
    only one way: The constructor takes one additional optional argument ``locales``
    expecting a list of language codes. If given, this instance is used to select
    the visible languages in all I18nFormFields of the form. If not given, all languages
    from ``settings.LANGUAGES`` will be displayed.

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nModelForm(BaseI18nModelForm, metaclass=ModelFormMetaclass):
    """
    This is a modified version of Django's ModelForm which differs from ModelForm in
    only one way: The constructor takes one additional optional argument ``locales``
    expecting a list of language codes. If given, this instance is used to select
    the visible languages in all I18nFormFields of the form. If not given, all languages
    from ``settings.LANGUAGES`` will be displayed.

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nFormSetMixin:

    def __init__(self, *args, **kwargs):
        self.locales = kwargs.pop('locales', None)
        super().__init__(*args, **kwargs)

    def _construct_form(self, i, **kwargs):
        kwargs['locales'] = self.locales
        return super()._construct_form(i, **kwargs)

    @property
    def empty_form(self):
        form = self.form(
            auto_id=self.auto_id,
            prefix=self.add_prefix('__prefix__'),
            empty_permitted=True,
            use_required_attribute=False,
            locales=self.locales
        )
        self.add_fields(form, None)
        return form

class I18nModelFormSet(I18nFormSetMixin, BaseModelFormSet):
    """
    This is equivalent to a normal BaseModelFormset, but cares for the special needs
    of I18nForms (see there for more information).

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nInlineFormSet(I18nFormSetMixin, BaseInlineFormSet):
    """
    This is equivalent to a normal BaseInlineFormset, but cares for the special needs
    of I18nForms (see there for more information).

    :param locales: A list of locales that should be displayed.
    """
    pass
<commit_msg>Fix empty formset forms with extra kwargs (#29)<commit_after>import copy
from typing import List, Union

from django import forms
from django.conf import settings
from django.forms import (
    BaseForm, BaseInlineFormSet, BaseModelForm, BaseModelFormSet,
)
from django.forms.forms import DeclarativeFieldsMetaclass
from django.forms.models import ModelFormMetaclass
from django.utils.safestring import mark_safe

from .strings import LazyI18nString

class I18nWidget(forms.MultiWidget):
    """
    The default form widget for I18nCharField and I18nTextField. It makes
    use of Django's MultiWidget mechanism and does some magic to save you
    time.
    """
    widget = forms.TextInput

    def __init__(self, locales: List[str], field: forms.Field, attrs=None):
        widgets = []
        self.locales = locales
        self.enabled_locales = locales
        self.field = field
        for lng in self.locales:
            a = copy.copy(attrs) or {}
            a['lang'] = lng
            widgets.append(self.widget(attrs=a))
        super().__init__(widgets, attrs)

    def decompress(self, value) -> List[Union[str, None]]:
        data = []
        first_enabled = None
        any_enabled_filled = False
        if not isinstance(value, LazyI18nString):
            value = LazyI18nString(value)
        for i, lng in enumerate(self.locales):
            dataline = (
                value.data[lng]
                if value is not None and (
                    isinstance(value.data, dict) or isinstance(value.data, LazyI18nString.LazyGettextProxy)
                ) and lng in value.data
                else None
            )
            if lng in self.enabled_locales:
                if not first_enabled:
                    first_enabled = i
                if dataline:
                    any_enabled_filled = True
            data.append(dataline)
        if value and not isinstance(value.data, dict) and not isinstance(value.data, LazyI18nString.LazyGettextProxy):
            data[first_enabled] = value.data
        elif value and not any_enabled_filled:
            data[first_enabled] = value.localize(self.enabled_locales[0])
        return data

    def render(self, name: str, value, attrs=None, renderer=None) -> str:
        if self.is_localized:
            for widget in self.widgets:
                widget.is_localized = self.is_localized
        # value is a list of values, each corresponding to a widget
        # in self.widgets.
        if not isinstance(value, list):
            value = self.decompress(value)
        output = []
        final_attrs = self.build_attrs(attrs or dict())
        id_ = final_attrs.get('id', None)
        for i, widget in enumerate(self.widgets):
            if self.locales[i] not in self.enabled_locales:
                continue
            try:
                widget_value = value[i]
            except IndexError:
                widget_value = None
            if id_:
                final_attrs = dict(
                    final_attrs,
                    id='%s_%s' % (id_, i),
                    title=self.locales[i]
                )
            output.append(widget.render(name + '_%s' % i, widget_value, final_attrs, renderer=renderer))
        return mark_safe(self.format_output(output))

    def format_output(self, rendered_widgets) -> str:
        return '
%s
' % (
            ' i18n-form-single-language' if len(rendered_widgets) <= 1 else '',
            ''.join(rendered_widgets),
        )

class I18nTextInput(I18nWidget):
    """
    The default form widget for I18nCharField. It makes use of Django's MultiWidget
    mechanism and does some magic to save you time.
    """
    widget = forms.TextInput

class I18nTextarea(I18nWidget):
    """
    The default form widget for I18nTextField. It makes use of Django's MultiWidget
    mechanism and does some magic to save you time.
    """
    widget = forms.Textarea

class I18nFormField(forms.MultiValueField):
    """
    The form field that is used by I18nCharField and I18nTextField. It makes use
    of Django's MultiValueField mechanism to create one sub-field per available
    language.

    It contains special treatment to make sure that a field marked as "required" is validated
    as "filled out correctly" if *at least one* translation is filled it. It is never required
    to fill in all of them. This has the drawback that the HTML property ``required`` is set on
    none of the fields as this would lead to irritating behaviour.

    :param locales: An iterable of locale codes that the widget should render a field for. If
                    omitted, fields will be rendered for all languages configured in
                    ``settings.LANGUAGES``.
    :param require_all_fields: A boolean, if set to True field requires all translations to be given.
    """

    def compress(self, data_list) -> LazyI18nString:
        locales = self.locales
        data = {}
        for i, value in enumerate(data_list):
            data[locales[i]] = value
        return LazyI18nString(data)

    def clean(self, value) -> LazyI18nString:
        if isinstance(value, LazyI18nString):
            # This happens e.g. if the field is disabled
            return value
        found = False
        found_all = True
        clean_data = []
        errors = []
        for i, field in enumerate(self.fields):
            try:
                field_value = value[i]
            except (IndexError, TypeError):
                field_value = None
            if field_value not in self.empty_values:
                found = True
            elif field.locale in self.widget.enabled_locales:
                found_all = False
            try:
                clean_data.append(field.clean(field_value))
            except forms.ValidationError as e:
                # Collect all validation errors in a single list, which we'll
                # raise at the end of clean(), rather than raising a single
                # exception for the first error we encounter. Skip duplicates.
                errors.extend(m for m in e.error_list if m not in errors)
        if errors:
            raise forms.ValidationError(errors)
        if self.one_required and not found:
            raise forms.ValidationError(self.error_messages['required'], code='required')
        if self.require_all_fields and not found_all:
            raise forms.ValidationError(self.error_messages['incomplete'], code='incomplete')

        out = self.compress(clean_data)
        self.validate(out)
        self.run_validators(out)
        return out

    def __init__(self, *args, **kwargs):
        fields = []
        defaults = {
            'widget': self.widget,
            'max_length': kwargs.pop('max_length', None),
        }
        self.locales = kwargs.pop('locales', [loc[0] for loc in settings.LANGUAGES])
        self.one_required = kwargs.get('required', True)
        require_all_fields = kwargs.pop('require_all_fields', False)
        kwargs['required'] = False
        kwargs['widget'] = kwargs['widget'](
            locales=self.locales, field=self, **kwargs.pop('widget_kwargs', {})
        )
        defaults.update(**kwargs)
        for lngcode in self.locales:
            defaults['label'] = '%s (%s)' % (defaults.get('label'), lngcode)
            field = forms.CharField(**defaults)
            field.locale = lngcode
            fields.append(field)
        super().__init__(
            fields=fields, require_all_fields=False, *args, **kwargs
        )
        self.require_all_fields = require_all_fields

class I18nFormMixin:
    def __init__(self, *args, **kwargs):
        locales = kwargs.pop('locales', None)
        super().__init__(*args, **kwargs)
        if locales:
            for k, field in self.fields.items():
                if isinstance(field, I18nFormField):
                    field.widget.enabled_locales = locales

class BaseI18nModelForm(I18nFormMixin, BaseModelForm):
    """
    This is a helperclass to construct an I18nModelForm.
    """
    pass

class BaseI18nForm(I18nFormMixin, BaseForm):
    """
    This is a helperclass to construct an I18nForm.
    """
    pass

class I18nForm(BaseI18nForm, metaclass=DeclarativeFieldsMetaclass):
    """
    This is a modified version of Django's Form which differs from Form in
    only one way: The constructor takes one additional optional argument ``locales``
    expecting a list of language codes. If given, this instance is used to select
    the visible languages in all I18nFormFields of the form. If not given, all languages
    from ``settings.LANGUAGES`` will be displayed.

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nModelForm(BaseI18nModelForm, metaclass=ModelFormMetaclass):
    """
    This is a modified version of Django's ModelForm which differs from ModelForm in
    only one way: The constructor takes one additional optional argument ``locales``
    expecting a list of language codes. If given, this instance is used to select
    the visible languages in all I18nFormFields of the form. If not given, all languages
    from ``settings.LANGUAGES`` will be displayed.

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nFormSetMixin:

    def __init__(self, *args, **kwargs):
        self.locales = kwargs.pop('locales', None)
        super().__init__(*args, **kwargs)

    def _construct_form(self, i, **kwargs):
        kwargs['locales'] = self.locales
        return super()._construct_form(i, **kwargs)

    @property
    def empty_form(self):
        form = self.form(
            auto_id=self.auto_id,
            prefix=self.add_prefix('__prefix__'),
            empty_permitted=True,
            use_required_attribute=False,
            locales=self.locales,
            **self.get_form_kwargs(None),
        )
        self.add_fields(form, None)
        return form

class I18nModelFormSet(I18nFormSetMixin, BaseModelFormSet):
    """
    This is equivalent to a normal BaseModelFormset, but cares for the special needs
    of I18nForms (see there for more information).

    :param locales: A list of locales that should be displayed.
    """
    pass

class I18nInlineFormSet(I18nFormSetMixin, BaseInlineFormSet):
    """
    This is equivalent to a normal BaseInlineFormset, but cares for the special needs
    of I18nForms (see there for more information).

    :param locales: A list of locales that should be displayed.
    """
    pass
<|endoftext|>


<commit_before>import itertools
import json
import uuid

from abc import ABCMeta
from functools import reduce
from operator import and_

import sqlalchemy as sa

from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.sql.elements import ClauseElement
from sqlalchemy import func

try:
    from modeltranslation import translator as dtrans
except ImportError:
    dtrans = None

from .decorators import method_connect_once, method_redis_once
from .. import utils, exceptions, aviews

CACHE_CATEGORY_COUNT = 'count'
CACHE_CATEGORY_SUM = 'aggregate:sum'

class MetaModel(ABCMeta):
    def __new__(mcls, name, bases, namespace):
        cls = ABCMeta.__new__(mcls, name, bases, namespace)
        name = utils.convert_class_name(name)
        cls.models[name] = cls
        return cls

class Model(dict, metaclass=MetaModel):
    models = {}
    app = None
    primary_key = 'id'
    validators = ()  # Validators for data
    update_validators = ()  # Validators to validate object and data before update
    fields_permanent = ()  # Fields need to be saved
    fields_readonly = ()
    fields_list = ()
    fields_one = None
    fields_localized = None

    @classmethod
    def factory(cls, app):
        return type(cls.__name__, (cls,), {'app': app})

    @classmethod
    def get_cache_key(cls, *args):
        parts = []
        if hasattr(cls.app, 'name'):
            parts.append(cls.app.name)
        else:
            parts.append(cls.app.__class__.__module__)
        parts.append(cls.__name__)
        parts.extend(args)
        return ':'.join(parts)

    def copy_object(self):
        cls = type(self)
        obj = cls(
            (n, v)
            for n, v in dict.items(self)
            if not isinstance(v, ClauseElement))
        return obj

    def pretty(self):
        return json.dumps(self, indent=3, cls=aviews.JsonEncoder)

    @property
    def pk(self):
        return self.get(self.primary_key)

    @pk.setter
    def pk(self, value):
        self[self.primary_key] = value

    def __getattr__(self, item):
        if item in self:
            return self[item]
        return getattr(super(), item)

    def __setattr__(self, key, value):
        if key == 'pk':
            key = self.primary_key
        self[key] = value

    @classmethod
    def _where(cls, args):
        if not args:
            raise ValueError('Where where?')
        elif isinstance(args[0], (int, str, uuid.UUID)):
            first, *tail = args
            args = [cls.table.c[cls.primary_key] == first]
            args.extend(tail)
        return reduce(and_, args),

    @classmethod
    def to_column(cls, fields):
        t = cls.table.c
        result = []
        for f in fields:
            if isinstance(f, str):
                # table.c returns an instance of ColumnCollection
                # and it has __getitem__ method to get column by it name.
                f = t[f]
            result.append(f)
        return result

    @classmethod
    def set_defaults(cls, data: dict):
        pass

    @classmethod
    async def _get_one(cls, *args, connection=None, fields=None):
        args = cls._where(args)
        if fields:
            fields = cls.to_column(fields)
        elif cls.fields_one:
            fields = cls.to_column(cls.fields_one)

        if fields:
            sql = sa.select(fields).select_from(cls.table)
        else:
            sql = cls.table.select()

        result = await connection.execute(sql.where(*args))
        return await result.first()

    @classmethod
    @method_connect_once
    async def get_one(cls, *args, connection=None, fields=None, silent=False):
        """
        Extract by id
        """
        r = await cls._get_one(*args, connection=connection, fields=fields)
        if r:
            return cls(**r)
        elif not silent:
            raise exceptions.NotFound()

    @method_connect_once
    async def load_fields(self, *fields, connection, force_update=False):
        fields = set(fields)
        if force_update is False:
            fields = fields - set(self)
        elif isinstance(force_update, (list, tuple)):
            fields = fields.union(force_update)

        if fields:
            r = await self._get_one(
                *self._where([self.pk]),
                connection=connection,
                fields=fields)
            dict.update(self, r)

    @classmethod
    @method_connect_once
    async def get_list(cls, *args, connection, fields=None,
                       offset=None, limit=None, sort=None,
                       select_from=None):
        """Extract list"""
        if fields:
            fields = cls.to_column(fields)
        elif cls.fields_list:
            fields = cls.to_column(cls.fields_list)

        if fields:
            sql = sa.select(fields).select_from(cls.table)
        else:
            sql = cls.table.select()

        for i in select_from or ():
            sql = sql.select_from(i)

        if args and args[0] is not None:
            sql = sql.where(reduce(and_, args))

        if offset is not None:
            sql = sql.offset(offset)

        if limit is not None:
            sql = sql.limit(limit)

        if isinstance(sort, str):
            sql = sql.order_by(sort)
        elif sort:
            sql = sql.order_by(*sort)

        result = await connection.execute(sql)
        l = []
        async for row in result:
            l.append(cls(**row))
        return l

    @classmethod
    @method_connect_once
    async def get_dict(cls, *where_and, connection=None,
                       fields=None, sort=None, **kwargs):
        where = []
        if where_and:
            if isinstance(where_and[0], (list, tuple, str, int)):
                v, *where_and = where_and
                kwargs[cls.primary_key] = v
        for k, v in kwargs.items():
            if isinstance(v, (list, tuple)):
                if v:
                    where.append(cls.table.c[k].in_(v))
            else:
                where.append(cls.table.c[k] == v)

        where.extend(where_and)
        if where:
            where = (reduce(and_, where),)
        else:
            where = ()
        if not fields:
            fields = None
        elif cls.primary_key not in fields:
            fields.append(cls.primary_key)
        l = await cls.get_list(
            *where, connection=connection,
            sort=sort, fields=fields)
        return {i.pk: i for i in l}

    @classmethod
    def get_table_from_django(cls, model, *jsonb, **field_type):
        """Deprecated, use @derive_from_django instead"""
        from .convert import convert_model
        for i in jsonb:
            field_type[i] = JSONB
        table, _ = convert_model(model, **field_type)
        return table

    @classmethod
    @method_connect_once
    async def _pg_scalar(cls, sql, connection=None):
        return await connection.scalar(sql)

    @classmethod
    @method_redis_once
    async def get_count(cls, *args, postfix=None, connection=None, redis=None, expire=180):
        """
        Extract query size
        """
        sql = cls.table.count()

        if args:
            sql = sql.where(reduce(and_, args))

        async def real_count():
            return await cls._pg_scalar(sql=sql, connection=connection)

        if expire == 0:
            return await real_count()

        if not postfix:
            postfix = _hash_stmt(sql)

        key = cls.get_cache_key(CACHE_CATEGORY_COUNT, postfix)

        count = await redis.get(key)
        if count is not None:
            return int(count)

        count = await real_count()
        await redis.set(key, count)
        await redis.expire(key, expire)

        return count

    @classmethod
    @method_connect_once
    @method_redis_once
    async def get_sum(cls, column, where, postfix=None, delay=0,
                      connection=None, redis=None):
        """Calculates sum"""
        sql = sa.select([func.sum(cls.table.c[column])]).where(where)

        if not postfix:
            postfix = _hash_stmt(sql)

        key = cls.get_cache_key(CACHE_CATEGORY_SUM, postfix)

        if delay:
            count = await redis.get(key)
            if count is not None:
                return int(count)

        count = await cls._pg_scalar(sql=sql, connection=connection)

        if count is None:
            count = 0
        elif delay:
            await redis.set(key, count)
            await redis.expire(key, delay)

        return count

    @classmethod
    @method_connect_once
    async def create(cls, *, connection, **kwargs):
        """Inserts new object"""
        pk = cls.table.c[cls.primary_key]
        cls.set_defaults(kwargs)
        uid = await connection.scalar(
            cls.table.insert().returning(pk).values(kwargs))
        kwargs[cls.primary_key] = uid
        return cls(**kwargs)

    @classmethod
    @method_connect_once
    async def create_many(cls, objects, connection=None):
        """Inserts many objects"""
        # aiopg doesn't support executemany so create object via cycle
        result = []
        for obj in objects:
            cls.set_defaults(obj)
            result.append(
                await cls.create(**obj, connection=connection)
            )
        return result

    @method_connect_once
    async def save(self, *, fields=None, connection):
        pk_field = self.table.c[self.primary_key]
        self.set_defaults(self)
        if self.primary_key in self:
            saved = await self._get_one(self.pk, connection=connection)
        else:
            saved = False
        if not saved:
            pk = await connection.scalar(
                self.table.insert().returning(pk_field).values(self))
            self[self.primary_key] = pk
            return pk
        if fields:
            fields = list(itertools.chain(fields, self.fields_permanent))
            values = {k: v for k, v in self.items()
                      if k in fields}
        elif self.fields_readonly:
            values = {k: v for k, v in self.items()
                      if k not in self.fields_readonly}
        else:
            values = self
        pk = await connection.scalar(
            self.table.update()
            .where(pk_field == self.pk)
            .returning(pk_field)
            .values(values)
        )
        assert self.pk == pk

        return pk

    @method_connect_once
    async def update_increment(self, connection=None, **kwargs):
        t = self.table

        dict_update = {
            t.c[field]: t.c[field] + value
            for field, value in kwargs.items()
        }

        await connection.execute(
            t.update().where(
                t.c[self.primary_key] == self.pk
            ).values(dict_update))

    @classmethod
    @method_connect_once
    async def update_fields(cls, where, connection=None, **kwargs):
        t = cls.table

        dict_update = {
            t.c[field]: value
            for field, value in kwargs.items()
        }

        await connection.execute(
            t.update().
            where(where).
            values(dict_update))

    @method_connect_once
    async def update_json(self, *args, connection=None, **kwargs):
        t = self.table
        if args:
            if len(args) > 1 and not kwargs:
                field, *path, value = args
            else:
                field, *path = args
                value = kwargs
            for p in reversed(path):
                value = {p: value}
            kwargs = {field: value}
        elif not kwargs:
            raise ValueError('Need args or kwargs')

        await connection.scalar(
            t.update().where(
                t.c[self.primary_key] == self.pk
            ).values(
                {
                    t.c[field]: sa.func.coalesce(
                        t.c[field], sa.cast({}, JSONB)
                    ) + sa.cast(value, JSONB)
                    for field, value in kwargs.items()
                }
            ).returning(t.c[self.primary_key]))

    @classmethod
    @method_connect_once
    async def delete_where(cls, *where, connection=None):
        t = cls.table

        where = cls._where(where)

        await connection.execute(
            t.delete().where(*where))

    @method_connect_once
    async def delete(self, connection=None):
        pk_field = self.table.c[self.primary_key]
        await connection.execute(self.table.delete().where(pk_field == self.pk))

    @classmethod
    @method_connect_once
    async def get_or_create(cls, *args, defaults=None, connection):
        pk_field = getattr(cls.table.c, cls.primary_key)
        if args:
            pass
        elif cls.primary_key in defaults:
            args = (defaults[cls.primary_key],)

        if args:
            saved = await cls._get_one(*args, connection=connection)
            if saved:
                return saved, False

        pk = await connection.scalar(
            cls.table.insert().returning(pk_field).values(defaults))
        obj = cls(**defaults)
        obj.pk = pk
        return obj, True

    @classmethod
    def validate(cls, data, to_class=True, default_validator=True):
        """Returns valid object or exception"""
        validators = cls.validators
        if not validators and default_validator:
            validators = [cls.default_validator]
        for validator in validators:
            data = validator(data)
        return cls(**data) if to_class else data

    @classmethod
    def default_validator(cls, data):
        return {f: data.get(f) for f in cls.table.columns.keys()}

    @method_connect_once
    async def validate_and_save(self, data, connection=None):
        """
        Method performs default validations, update validations and save object.
        """
        # Validate data using user defined validators
        data = self.validate(data, to_class=False, default_validator=False)
        for v in self.update_validators:
            data = v(self, data)
        # Do not allow object update for empty data to avoid extra save.
        if data:
            self.update(data)
            return await self.save(fields=data.keys(), connection=connection)

    @staticmethod
    def get_fields_localized_from_django(django_model):
        if dtrans:
            try:
                return list(dtrans.translator.get_options_for_model(django_model).fields.keys())
            except dtrans.NotRegistered:
                pass

    @classmethod
    def localize(cls, obj, locale):
        if not cls.fields_localized:
            return
        for field in cls.fields_localized:
            if field not in obj:
                continue
            value = obj.get('{}_{}'.format(field, locale))
            if value:
                obj[field] = value

def _hash_stmt(stmt):
    compiled = stmt.compile()
    msg = compiled.string + repr(compiled.params)
    return utils.get_hash(msg)
<commit_msg>Model work on asyncpg/asyncpgsa<commit_after>import itertools
import json
import uuid

from abc import ABCMeta
from functools import reduce
from operator import and_

import sqlalchemy as sa

from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.sql.elements import ClauseElement
from sqlalchemy import func

try:
    from modeltranslation import translator as dtrans
except ImportError:
    dtrans = None

from .decorators import method_connect_once, method_redis_once
from .. import utils, exceptions, aviews

CACHE_CATEGORY_COUNT = 'count'
CACHE_CATEGORY_SUM = 'aggregate:sum'

class MetaModel(ABCMeta):
    def __new__(mcls, name, bases, namespace):
        cls = ABCMeta.__new__(mcls, name, bases, namespace)
        name = utils.convert_class_name(name)
        cls.models[name] = cls
        return cls

class Model(dict, metaclass=MetaModel):
    models = {}
    app = None
    primary_key = 'id'
    validators = ()  # Validators for data
    update_validators = ()  # Validators to validate object and data before update
    fields_permanent = ()  # Fields need to be saved
    fields_readonly = ()
    fields_list = ()
    fields_one = None
    fields_localized = None

    @classmethod
    def factory(cls, app):
        return type(cls.__name__, (cls,), {'app': app})

    @classmethod
    def get_cache_key(cls, *args):
        parts = []
        if hasattr(cls.app, 'name'):
            parts.append(cls.app.name)
        else:
            parts.append(cls.app.__class__.__module__)
        parts.append(cls.__name__)
        parts.extend(args)
        return ':'.join(parts)

    def copy_object(self):
        cls = type(self)
        obj = cls(
            (n, v)
            for n, v in dict.items(self)
            if not isinstance(v, ClauseElement))
        return obj

    def pretty(self):
        return json.dumps(self, indent=3, cls=aviews.JsonEncoder)

    @property
    def pk(self):
        return self.get(self.primary_key)

    @pk.setter
    def pk(self, value):
        self[self.primary_key] = value

    def __getattr__(self, item):
        if item in self:
            return self[item]
        return getattr(super(), item)

    def __setattr__(self, key, value):
        if key == 'pk':
            key = self.primary_key
        self[key] = value

    @classmethod
    def _where(cls, args):
        if not args:
            raise ValueError('Where where?')
        elif isinstance(args[0], (int, str, uuid.UUID)):
            first, *tail = args
            args = [cls.table.c[cls.primary_key] == first]
            args.extend(tail)
        return reduce(and_, args),

    @classmethod
    def to_column(cls, fields):
        t = cls.table.c
        result = []
        for f in fields:
            if isinstance(f, str):
                # table.c returns an instance of ColumnCollection
                # and it has __getitem__ method to get column by it name.
                f = t[f]
            result.append(f)
        return result

    @classmethod
    def set_defaults(cls, data: dict):
        pass

    @classmethod
    async def _get_one(cls, *args, connection=None, fields=None):
        args = cls._where(args)
        if fields:
            fields = cls.to_column(fields)
        elif cls.fields_one:
            fields = cls.to_column(cls.fields_one)

        if fields:
            sql = sa.select(fields).select_from(cls.table)
        else:
            sql = cls.table.select()

        return await connection.fetchrow(sql.where(*args))

    @classmethod
    @method_connect_once
    async def get_one(cls, *args, connection=None, fields=None, silent=False):
        """
        Extract by id
        """
        r = await cls._get_one(*args, connection=connection, fields=fields)
        if r:
            return cls(**r)
        elif not silent:
            raise exceptions.NotFound()

    @method_connect_once
    async def load_fields(self, *fields, connection, force_update=False):
        fields = set(fields)
        if force_update is False:
            fields = fields - set(self)
        elif isinstance(force_update, (list, tuple)):
            fields = fields.union(force_update)

        if fields:
            r = await self._get_one(
                *self._where([self.pk]),
                connection=connection,
                fields=fields)
            dict.update(self, r)

    @classmethod
    @method_connect_once
    async def get_list(cls, *args, connection, fields=None,
                       offset=None, limit=None, sort=None,
                       select_from=None):
        """Extract list"""
        if fields:
            fields = cls.to_column(fields)
        elif cls.fields_list:
            fields = cls.to_column(cls.fields_list)

        if fields:
            sql = sa.select(fields).select_from(cls.table)
        else:
            sql = cls.table.select()

        for i in select_from or ():
            sql = sql.select_from(i)

        if args and args[0] is not None:
            sql = sql.where(reduce(and_, args))

        if offset is not None:
            sql = sql.offset(offset)

        if limit is not None:
            sql = sql.limit(limit)

        if isinstance(sort, str):
            sql = sql.order_by(sort)
        elif sort:
            sql = sql.order_by(*sort)

        result = await connection.fetch(sql)
        return [cls(**row) for row in result]

    @classmethod
    @method_connect_once
    async def get_dict(cls, *where_and, connection=None,
                       fields=None, sort=None, **kwargs):
        where = []
        if where_and:
            if isinstance(where_and[0], (list, tuple, str, int)):
                v, *where_and = where_and
                kwargs[cls.primary_key] = v
        for k, v in kwargs.items():
            if isinstance(v, (list, tuple)):
                if v:
                    where.append(cls.table.c[k].in_(v))
            else:
                where.append(cls.table.c[k] == v)

        where.extend(where_and)
        if where:
            where = (reduce(and_, where),)
        else:
            where = ()
        if not fields:
            fields = None
        elif cls.primary_key not in fields:
            fields.append(cls.primary_key)
        l = await cls.get_list(
            *where, connection=connection,
            sort=sort, fields=fields)
        return {i.pk: i for i in l}

    @classmethod
    def get_table_from_django(cls, model, *jsonb, **field_type):
        """Deprecated, use @derive_from_django instead"""
        from .convert import convert_model
        for i in jsonb:
            field_type[i] = JSONB
        table, _ = convert_model(model, **field_type)
        return table

    @classmethod
    @method_connect_once
    async def _pg_scalar(cls, sql, connection=None):
        return await connection.fetchval(sql)

    @classmethod
    @method_redis_once
    async def get_count(cls, *args, postfix=None, connection=None, redis=None, expire=180):
        """
        Extract query size
        """
        sql = cls.table.count()

        if args:
            sql = sql.where(reduce(and_, args))

        async def real_count():
            return await cls._pg_scalar(sql=sql, connection=connection)

        if expire == 0:
            return await real_count()

        if not postfix:
            postfix = _hash_stmt(sql)

        key = cls.get_cache_key(CACHE_CATEGORY_COUNT, postfix)

        count = await redis.get(key)
        if count is not None:
            return int(count)

        count = await real_count()
        await redis.set(key, count)
        await redis.expire(key, expire)

        return count

    @classmethod
    @method_connect_once
    @method_redis_once
    async def get_sum(cls, column, where, postfix=None, delay=0,
                      connection=None, redis=None):
        """Calculates sum"""
        sql = sa.select([func.sum(cls.table.c[column])]).where(where)

        if not postfix:
            postfix = _hash_stmt(sql)

        key = cls.get_cache_key(CACHE_CATEGORY_SUM, postfix)

        if delay:
            count = await redis.get(key)
            if count is not None:
                return int(count)

        count = await cls._pg_scalar(sql=sql, connection=connection)

        if count is None:
            count = 0
        elif delay:
            await redis.set(key, count)
            await redis.expire(key, delay)

        return count

    @classmethod
    @method_connect_once
    async def create(cls, *, connection, **kwargs):
        """Inserts new object"""
        pk = cls.table.c[cls.primary_key]
        cls.set_defaults(kwargs)
        uid = await connection.fetchval(
            cls.table.insert().returning(pk).values(kwargs))
        kwargs[cls.primary_key] = uid
        return cls(**kwargs)

    @classmethod
    @method_connect_once
    async def create_many(cls, objects, connection=None, returning=True):
        """Inserts many objects"""
        pk = cls.table.c[cls.primary_key]
        sql = cls.table.insert()
        if returning:
            sql = sql.returning(pk)
        for obj in objects:
            cls.set_defaults(obj)
        sql = sql.values(objects)
        if returning:
            result = await connection.fetch(sql)
            for pk, obj in zip(result, objects):
                obj[cls.primary_key] = pk[cls.primary_key]
            return [cls(**obj) for obj in objects]
        else:
            await connection.execute(sql)

    @method_connect_once
    async def save(self, *, fields=None, connection):
        pk_field = self.table.c[self.primary_key]
        self.set_defaults(self)
        if self.primary_key in self:
            saved = await self._get_one(self.pk, connection=connection)
        else:
            saved = False
        if not saved:
            pk = await connection.fetchval(
                self.table.insert().returning(pk_field).values(self))
            self[self.primary_key] = pk
            return pk
        if fields:
            fields = list(itertools.chain(fields, self.fields_permanent))
            values = {k: v for k, v in self.items()
                      if k in fields}
        elif self.fields_readonly:
            values = {k: v for k, v in self.items()
                      if k not in self.fields_readonly}
        else:
            values = self
        pk = await connection.fetchval(
            self.table.update()
            .where(pk_field == self.pk)
            .returning(pk_field)
            .values(values)
        )
        assert self.pk == pk

        return pk

    @method_connect_once
    async def update_increment(self, connection=None, **kwargs):
        t = self.table

        dict_update = {
            t.c[field]: t.c[field] + value
            for field, value in kwargs.items()
        }

        await connection.execute(
            t.update().where(
                t.c[self.primary_key] == self.pk
            ).values(dict_update))

    @classmethod
    @method_connect_once
    async def update_fields(cls, where, connection=None, **kwargs):
        t = cls.table

        dict_update = {
            t.c[field]: value
            for field, value in kwargs.items()
        }

        await connection.execute(
            t.update().
            where(where).
            values(dict_update))

    @method_connect_once
    async def update_json(self, *args, connection=None, **kwargs):
        t = self.table
        if args:
            if len(args) > 1 and not kwargs:
                field, *path, value = args
            else:
                field, *path = args
                value = kwargs
            for p in reversed(path):
                value = {p: value}
            kwargs = {field: value}
        elif not kwargs:
            raise ValueError('Need args or kwargs')

        await connection.fetchval(
            t.update().where(
                t.c[self.primary_key] == self.pk
            ).values(
                {
                    t.c[field]: sa.func.coalesce(
                        t.c[field], sa.cast({}, JSONB)
                    ) + sa.cast(value, JSONB)
                    for field, value in kwargs.items()
                }
            ).returning(t.c[self.primary_key]))

    @classmethod
    @method_connect_once
    async def delete_where(cls, *where, connection=None):
        t = cls.table

        where = cls._where(where)

        await connection.fetchval(
            t.delete().where(*where))

    @method_connect_once
    async def delete(self, connection=None):
        pk_field = self.table.c[self.primary_key]
        await connection.fetchval(self.table.delete().where(pk_field == self.pk))

    @classmethod
    @method_connect_once
    async def get_or_create(cls, *args, defaults=None, connection):
        pk_field = getattr(cls.table.c, cls.primary_key)
        if args:
            pass
        elif cls.primary_key in defaults:
            args = (defaults[cls.primary_key],)

        if args:
            saved = await cls._get_one(*args, connection=connection)
            if saved:
                return saved, False

        pk = await connection.fetchval(
            cls.table.insert().returning(pk_field).values(defaults))
        obj = cls(**defaults)
        obj.pk = pk
        return obj, True

    @classmethod
    def validate(cls, data, to_class=True, default_validator=True):
        """Returns valid object or exception"""
        validators = cls.validators
        if not validators and default_validator:
            validators = [cls.default_validator]
        for validator in validators:
            data = validator(data)
        return cls(**data) if to_class else data

    @classmethod
    def default_validator(cls, data):
        return {f: data.get(f) for f in cls.table.columns.keys()}

    @method_connect_once
    async def validate_and_save(self, data, connection=None):
        """
        Method performs default validations, update validations and save object.
        """
        # Validate data using user defined validators
        data = self.validate(data, to_class=False, default_validator=False)
        for v in self.update_validators:
            data = v(self, data)
        # Do not allow object update for empty data to avoid extra save.
        if data:
            self.update(data)
            return await self.save(fields=data.keys(), connection=connection)

    @staticmethod
    def get_fields_localized_from_django(django_model):
        if dtrans:
            try:
                return list(dtrans.translator.get_options_for_model(django_model).fields.keys())
            except dtrans.NotRegistered:
                pass

    @classmethod
    def localize(cls, obj, locale):
        if not cls.fields_localized:
            return
        for field in cls.fields_localized:
            if field not in obj:
                continue
            value = obj.get('{}_{}'.format(field, locale))
            if value:
                obj[field] = value

def _hash_stmt(stmt):
    compiled = stmt.compile()
    msg = compiled.string + repr(compiled.params)
    return utils.get_hash(msg)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import collections
import contextlib
import functools
import itertools
import os
import shutil
import tempfile
import unittest
import urllib.parse

import mutagen
import requests

import sacad.recurse as recurse

try:
  redirect_stdout = contextlib.redirect_stdout
except AttributeError:
  # contextlib.redirect_stdout is not available (Python 3.3), build our own
  import sys
  @contextlib.contextmanager
  def redirect_stdout(s):
    original_stdout, sys.stdout = sys.stdout, s
    try:
      yield
    finally:
      sys.stdout = original_stdout

def download(url, filepath):
  cache_dir = os.getenv("TEST_DL_CACHE_DIR")
  if cache_dir is not None:
    os.makedirs(cache_dir, exist_ok=True)
    cache_filepath = os.path.join(cache_dir,
                                  os.path.basename(urllib.parse.urlsplit(url).path))
    if os.path.isfile(cache_filepath):
      shutil.copyfile(cache_filepath, filepath)
      return
  with contextlib.closing(requests.get(url, stream=True)) as response:
    response.raise_for_status()
    with open(filepath, "wb") as f:
      for chunk in response.iter_content(2 ** 14):
        f.write(chunk)
  if cache_dir is not None:
    shutil.copyfile(filepath, cache_filepath)

class TestRecursive(unittest.TestCase):

  @classmethod
  def setUpClass(cls):
    cls.temp_dir = tempfile.TemporaryDirectory()
    cls.album1_dir = os.path.join(cls.temp_dir.name, "album1")
    os.mkdir(cls.album1_dir)
    url = "https://upload.wikimedia.org/wikipedia/en/4/45/ACDC_-_Back_In_Black-sample.ogg"
    filepath1 = os.path.join(cls.album1_dir, "2 track.ogg")
    download(url, filepath1)
    mf = mutagen.File(filepath1)
    mf["artist"] = "ARTIST1"
    mf["album"] = "ALBUM1"
    mf.save()

    cls.album2_dir = os.path.join(cls.temp_dir.name, "album2")
    os.mkdir(cls.album2_dir)
    filepath2 = os.path.join(cls.album2_dir, "1.dat")
    with open(filepath2, "wb") as f:
      f.write(b"\x00" * 8)
    filepath3 = os.path.join(cls.album2_dir, "2 track.ogg")
    shutil.copyfile(filepath1, filepath3)
    mf = mutagen.File(filepath3)
    mf["artist"] = "ARTIST2"
    mf["album"] = "ALBUM2"
    mf.save()

    cls.album3_dir = os.path.join(cls.temp_dir.name, "album3")
    os.mkdir(cls.album3_dir)
    url = "https://allthingsaudio.wikispaces.com/file/view/Shuffle%20for%20K.M.mp3/139190697/Shuffle%20for%20K.M.mp3"
    download(url, os.path.join(cls.album3_dir, "1 track.mp3"))

    cls.album4_dir = os.path.join(cls.temp_dir.name, "album4")
    os.mkdir(cls.album4_dir)
    url = "https://auphonic.com/media/audio-examples/01.auphonic-demo-unprocessed.m4a"
    download(url, os.path.join(cls.album4_dir, "1 track.m4a"))

    cls.not_album_dir = os.path.join(cls.temp_dir.name, "not an album")
    os.mkdir(cls.not_album_dir)
    shutil.copyfile(filepath2, os.path.join(cls.not_album_dir, "a.dat"))

    cls.invalid_album_dir = os.path.join(cls.temp_dir.name, "invalid album")
    os.mkdir(cls.invalid_album_dir)
    shutil.copyfile(filepath2, os.path.join(cls.invalid_album_dir, "2 track.ogg"))
    filepath4 = os.path.join(cls.invalid_album_dir, "3 track.ogg")
    shutil.copyfile(filepath1, filepath4)
    mf = mutagen.File(filepath4)
    del mf["album"]
    mf.save()

  @classmethod
  def tearDownClass(cls):
    cls.temp_dir.cleanup()

  def test_analyze_lib(self):
    with open(os.devnull, "wt") as dn, redirect_stdout(dn):
      work = recurse.analyze_lib(__class__.temp_dir.name, "a.jpg")
      self.assertEqual(len(work), 4)
      self.assertIn(__class__.album1_dir, work)
      self.assertEqual(work[__class__.album1_dir], ("ARTIST1", "ALBUM1"))
      self.assertIn(__class__.album2_dir, work)
      self.assertEqual(work[__class__.album2_dir], ("ARTIST2", "ALBUM2"))
      self.assertIn(__class__.album3_dir, work)
      self.assertEqual(work[__class__.album3_dir], ("jpfmband", "Paris S.F""))
      self.assertIn(__class__.album4_dir, work)
      self.assertEqual(work[__class__.album4_dir], ("Auphonic", "Auphonic Demonstration"))

      work = recurse.analyze_lib(__class__.temp_dir.name, "1.dat")
      self.assertEqual(len(work), 3)
      self.assertIn(__class__.album1_dir, work)
      self.assertEqual(work[__class__.album1_dir], ("ARTIST1", "ALBUM1"))
      self.assertIn(__class__.album3_dir, work)
      self.assertEqual(work[__class__.album3_dir], ("jpfmband", "Paris S.F"))
      self.assertIn(__class__.album4_dir, work)
      self.assertEqual(work[__class__.album4_dir], ("Auphonic", "Auphonic Demonstration"))

  def test_get_metadata(self):
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album1_dir),
                                              os.listdir(__class__.album1_dir))),
                     ("ARTIST1", "ALBUM1"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album2_dir),
                                              os.listdir(__class__.album2_dir))),
                     ("ARTIST2", "ALBUM2"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album3_dir),
                                              os.listdir(__class__.album3_dir))),
                     ("jpfmband", "Paris S.F"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album4_dir),
                                              os.listdir(__class__.album4_dir))),
                     ("Auphonic", "Auphonic Demonstration"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.not_album_dir),
                                              os.listdir(__class__.not_album_dir))),
                     (None, None))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.invalid_album_dir),
                                              os.listdir(__class__.invalid_album_dir))),
                     ("ARTIST1", None))

  def test_analyze_dir(self):
    scrobbler = itertools.cycle("|/-\\")
    with open(os.devnull, "wt") as dn, redirect_stdout(dn):
      stats = collections.defaultdict(int)
      failed_dirs = []
      r = recurse.analyze_dir(stats,
                              __class__.album1_dir,
                              os.listdir(__class__.album1_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 1)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], ("ARTIST1", "ALBUM1"))

      stats.clear()
      r = recurse.analyze_dir(stats,
                              __class__.album2_dir,
                              os.listdir(__class__.album2_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], ("ARTIST2", "ALBUM2"))
      stats.clear()
      r = recurse.analyze_dir(stats,
                              __class__.album2_dir,
                              os.listdir(__class__.album2_dir),
                              "1.dat",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertNotIn("missing covers", stats)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], (None, None))

      stats.clear()
      failed_dirs.clear()
      r = recurse.analyze_dir(stats,
                              __class__.not_album_dir,
                              os.listdir(__class__.not_album_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 1)
      self.assertNotIn("albums", stats)
      self.assertNotIn("missing covers", stats)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], (None, None))

      stats.clear()
      failed_dirs.clear()
      r = recurse.analyze_dir(stats,
                              __class__.invalid_album_dir,
                              os.listdir(__class__.invalid_album_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertIn("errors", stats)
      self.assertEqual(stats["errors"], 1)
      self.assertSequenceEqual(failed_dirs, (__class__.invalid_album_dir,))
      self.assertEqual(r[0], ("ARTIST1", None))

if __name__ == "__main__":
  # run tests
  unittest.main()
<commit_msg>Fix tests - 2<commit_after>#!/usr/bin/env python3

import collections
import contextlib
import functools
import itertools
import os
import shutil
import tempfile
import unittest
import urllib.parse

import mutagen
import requests

import sacad.recurse as recurse

try:
  redirect_stdout = contextlib.redirect_stdout
except AttributeError:
  # contextlib.redirect_stdout is not available (Python 3.3), build our own
  import sys
  @contextlib.contextmanager
  def redirect_stdout(s):
    original_stdout, sys.stdout = sys.stdout, s
    try:
      yield
    finally:
      sys.stdout = original_stdout

def download(url, filepath):
  cache_dir = os.getenv("TEST_DL_CACHE_DIR")
  if cache_dir is not None:
    os.makedirs(cache_dir, exist_ok=True)
    cache_filepath = os.path.join(cache_dir,
                                  os.path.basename(urllib.parse.urlsplit(url).path))
    if os.path.isfile(cache_filepath):
      shutil.copyfile(cache_filepath, filepath)
      return
  with contextlib.closing(requests.get(url, stream=True)) as response:
    response.raise_for_status()
    with open(filepath, "wb") as f:
      for chunk in response.iter_content(2 ** 14):
        f.write(chunk)
  if cache_dir is not None:
    shutil.copyfile(filepath, cache_filepath)

class TestRecursive(unittest.TestCase):

  @classmethod
  def setUpClass(cls):
    cls.temp_dir = tempfile.TemporaryDirectory()
    cls.album1_dir = os.path.join(cls.temp_dir.name, "album1")
    os.mkdir(cls.album1_dir)
    url = "https://upload.wikimedia.org/wikipedia/en/4/45/ACDC_-_Back_In_Black-sample.ogg"
    filepath1 = os.path.join(cls.album1_dir, "2 track.ogg")
    download(url, filepath1)
    mf = mutagen.File(filepath1)
    mf["artist"] = "ARTIST1"
    mf["album"] = "ALBUM1"
    mf.save()

    cls.album2_dir = os.path.join(cls.temp_dir.name, "album2")
    os.mkdir(cls.album2_dir)
    filepath2 = os.path.join(cls.album2_dir, "1.dat")
    with open(filepath2, "wb") as f:
      f.write(b"\x00" * 8)
    filepath3 = os.path.join(cls.album2_dir, "2 track.ogg")
    shutil.copyfile(filepath1, filepath3)
    mf = mutagen.File(filepath3)
    mf["artist"] = "ARTIST2"
    mf["album"] = "ALBUM2"
    mf.save()

    cls.album3_dir = os.path.join(cls.temp_dir.name, "album3")
    os.mkdir(cls.album3_dir)
    url = "https://allthingsaudio.wikispaces.com/file/view/Shuffle%20for%20K.M.mp3/139190697/Shuffle%20for%20K.M.mp3"
    download(url, os.path.join(cls.album3_dir, "1 track.mp3"))

    cls.album4_dir = os.path.join(cls.temp_dir.name, "album4")
    os.mkdir(cls.album4_dir)
    url = "https://auphonic.com/media/audio-examples/01.auphonic-demo-unprocessed.m4a"
    download(url, os.path.join(cls.album4_dir, "1 track.m4a"))

    cls.not_album_dir = os.path.join(cls.temp_dir.name, "not an album")
    os.mkdir(cls.not_album_dir)
    shutil.copyfile(filepath2, os.path.join(cls.not_album_dir, "a.dat"))

    cls.invalid_album_dir = os.path.join(cls.temp_dir.name, "invalid album")
    os.mkdir(cls.invalid_album_dir)
    shutil.copyfile(filepath2, os.path.join(cls.invalid_album_dir, "2 track.ogg"))
    filepath4 = os.path.join(cls.invalid_album_dir, "3 track.ogg")
    shutil.copyfile(filepath1, filepath4)
    mf = mutagen.File(filepath4)
    del mf["album"]
    mf.save()

  @classmethod
  def tearDownClass(cls):
    cls.temp_dir.cleanup()

  def test_analyze_lib(self):
    with open(os.devnull, "wt") as dn, redirect_stdout(dn):
      work = recurse.analyze_lib(__class__.temp_dir.name, "a.jpg")
      self.assertEqual(len(work), 4)
      self.assertIn(__class__.album1_dir, work)
      self.assertEqual(work[__class__.album1_dir], ("ARTIST1", "ALBUM1"))
      self.assertIn(__class__.album2_dir, work)
      self.assertEqual(work[__class__.album2_dir], ("ARTIST2", "ALBUM2"))
      self.assertIn(__class__.album3_dir, work)
      self.assertEqual(work[__class__.album3_dir], ("jpfmband", "Paris S.F"))
      self.assertIn(__class__.album4_dir, work)
      self.assertEqual(work[__class__.album4_dir], ("Auphonic", "Auphonic Demonstration"))

      work = recurse.analyze_lib(__class__.temp_dir.name, "1.dat")
      self.assertEqual(len(work), 3)
      self.assertIn(__class__.album1_dir, work)
      self.assertEqual(work[__class__.album1_dir], ("ARTIST1", "ALBUM1"))
      self.assertIn(__class__.album3_dir, work)
      self.assertEqual(work[__class__.album3_dir], ("jpfmband", "Paris S.F"))
      self.assertIn(__class__.album4_dir, work)
      self.assertEqual(work[__class__.album4_dir], ("Auphonic", "Auphonic Demonstration"))

  def test_get_metadata(self):
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album1_dir),
                                              os.listdir(__class__.album1_dir))),
                     ("ARTIST1", "ALBUM1"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album2_dir),
                                              os.listdir(__class__.album2_dir))),
                     ("ARTIST2", "ALBUM2"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album3_dir),
                                              os.listdir(__class__.album3_dir))),
                     ("jpfmband", "Paris S.F"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.album4_dir),
                                              os.listdir(__class__.album4_dir))),
                     ("Auphonic", "Auphonic Demonstration"))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.not_album_dir),
                                              os.listdir(__class__.not_album_dir))),
                     (None, None))
    self.assertEqual(recurse.get_metadata(map(functools.partial(os.path.join,
                                                                __class__.invalid_album_dir),
                                              os.listdir(__class__.invalid_album_dir))),
                     ("ARTIST1", None))

  def test_analyze_dir(self):
    scrobbler = itertools.cycle("|/-\\")
    with open(os.devnull, "wt") as dn, redirect_stdout(dn):
      stats = collections.defaultdict(int)
      failed_dirs = []
      r = recurse.analyze_dir(stats,
                              __class__.album1_dir,
                              os.listdir(__class__.album1_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 1)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], ("ARTIST1", "ALBUM1"))

      stats.clear()
      r = recurse.analyze_dir(stats,
                              __class__.album2_dir,
                              os.listdir(__class__.album2_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], ("ARTIST2", "ALBUM2"))
      stats.clear()
      r = recurse.analyze_dir(stats,
                              __class__.album2_dir,
                              os.listdir(__class__.album2_dir),
                              "1.dat",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertNotIn("missing covers", stats)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], (None, None))

      stats.clear()
      failed_dirs.clear()
      r = recurse.analyze_dir(stats,
                              __class__.not_album_dir,
                              os.listdir(__class__.not_album_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 1)
      self.assertNotIn("albums", stats)
      self.assertNotIn("missing covers", stats)
      self.assertNotIn("errors", stats)
      self.assertEqual(len(failed_dirs), 0)
      self.assertEqual(r[0], (None, None))

      stats.clear()
      failed_dirs.clear()
      r = recurse.analyze_dir(stats,
                              __class__.invalid_album_dir,
                              os.listdir(__class__.invalid_album_dir),
                              "1.jpg",
                              failed_dirs,
                              scrobbler,
                              0)
      self.assertIn("files", stats)
      self.assertEqual(stats["files"], 2)
      self.assertIn("albums", stats)
      self.assertEqual(stats["albums"], 1)
      self.assertIn("missing covers", stats)
      self.assertEqual(stats["missing covers"], 1)
      self.assertIn("errors", stats)
      self.assertEqual(stats["errors"], 1)
      self.assertSequenceEqual(failed_dirs, (__class__.invalid_album_dir,))
      self.assertEqual(r[0], ("ARTIST1", None))

if __name__ == "__main__":
  # run tests
  unittest.main()
<|endoftext|>


<commit_before>from i3pystatus import IntervalModule
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_vertical_bar

class Sensor:
    """
    Simple class representing a CPU temperature sensor.
    """

    def __init__(self, name, current, maximum, critical):
        self.name = name.replace(' ', '_')
        self.current = int(current)
        self.maximum = int(maximum) if maximum else int(critical)
        self.critical = int(critical)

    def __repr__(self):
        return "Sensor(name='{}', current={}, maximum={}, critical={})".format(
            self.name,
            self.current,
            self.maximum,
            self.critical
        )

    def is_warning(self):
        return self.current > self.maximum

    def is_critical(self):
        return self.current > self.critical

def get_sensors():
    """ Detect and return a list of Sensor objects """
    import sensors
    found_sensors = list()

    def get_subfeature_value(feature, subfeature_type):
        subfeature = chip.get_subfeature(feature, subfeature_type)
        if subfeature:
            return chip.get_value(subfeature.number)

    for chip in sensors.get_detected_chips():
        for feature in chip.get_features():
            if feature.type == sensors.FEATURE_TEMP:
                try:
                    name = chip.get_label(feature)
                    max = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_MAX)
                    current = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_INPUT)
                    critical = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_CRIT)
                    if critical:
                        found_sensors.append(Sensor(name=name, current=current, maximum=max, critical=critical))
                except sensors.SensorsException:
                    continue
    return found_sensors

class Temperature(IntervalModule, ColorRangeModule):
    """
    Shows CPU temperature of Intel processors.

    AMD is currently not supported as they can only report a relative temperature, which is pretty useless.

    Requires `colour` module from PyPi

    .. rubric:: Modes of operation

    If lm_sensors_enabled is set to False, the module operates in default mode. This means that:
        * only the {temp} formatter is available
        * alert_temp is honored

    If lm_sensors_enabled is set to True, the module operates in lm_sensors mode. This means that:
        * pysensors must be installed (https://github.com/bastienleonard/pysensors)
        * CPU sensors are discovered dynamically (supporting a sensor per core and multiple CPUs)
        * alert_temp is ignored. The warning or critical values reported by the sensor are used instead (see urgent_on)

    .. rubric:: lm_sensors installation

    In order to take advantage of the lm_sensors library and tools, it must first be installed and configured.

    On Arch this is as simple as:

    .. code-block:: bash

        pacman -S lm_sensors

    On Ubuntu:

    .. code-block:: bash

        sudo apt-get update && sudo apt-get install lm-sensors libsensors4-dev

    The Arch Wiki has a good page on the library - https://wiki.archlinux.org/index.php/lm_sensors

    .. rubric:: lm_sensors_mode formatters

    When ``lm_sensors_enabled`` is True the formatters are created dynamically. In order to discover the formatters that
    are available, it is best to run the sensors command:

    .. code-block:: bash

        ⇒ sensors
        coretemp-isa-0000
        Adapter: ISA adapter
        Physical id 0:  +48.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 0:         +48.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 1:         +46.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 2:         +43.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 3:         +47.0°C  (high = +80.0°C, crit = +99.0°C)

    The module replaces spaces in sensor names with underscores, therefore from the above output we can
    identify the following sensor formatters:

    * Physical_id_0
    * Core_0
    * Core_1
    * Core_2
    * Core_3

    For each sensor a vertical bar is also generated. In this example we would also have the following bars:

    * Physical_id_0_bar
    * Core_0_bar
    * Core_1_bar
    * Core_2_bar
    * Core_3_bar

    Thus, this format string would be valid: "{Physical_id_0}°C {Core_0_bar}{Core_1_bar}{Core_2_bar}{Core_3_bar}"

    .. rubric:: Pango Markup and lm_sensors_mode

    When Pango Markup is enabled and ``dynamic_color`` is True, each sensor's formatter color is displayed independently.
    The color is determined by the proximity of the sensors current value to it's critical value.

    .. rubric:: Example Configuration

    Here is an example configuration based on the sensor values discovered above:

    .. code-block:: python

        status.register("temp",
                        format="{Physical_id_0}°C {Core_0_bar}{Core_1_bar}{Core_2_bar}{Core_3_bar}",
                        hints={"markup": "pango"},
                        lm_sensors_enabled=True,
                        dynamic_color=True)
    """

    settings = (
        ("format",
         "format string used for output. {temp} is the temperature in degrees celsius"),
        ('display_if', 'snippet that gets evaluated. if true, displays the module output'),
        ('lm_sensors_enabled', 'whether or not lm_sensors should be used for obtaining CPU temperature information'),
        ('urgent_on', 'whether to flag as urgent when temperature exceeds urgent value or critical value '
                      '(requires lm_sensors_enabled)'),
        ('dynamic_color', 'whether to set the color dynamically (overrides alert_color)'),
        "color",
        "file",
        "alert_temp",
        "alert_color",
    )
    format = "{temp} °C"
    color = "#FFFFFF"
    file = "/sys/class/thermal/thermal_zone0/temp"
    alert_temp = 90
    alert_color = "#FF0000"
    display_if = 'True'

    lm_sensors_enabled = False
    dynamic_color = False
    urgent_on = 'warning'

    def init(self):
        self.pango_enabled = self.hints.get("markup", False) and self.hints["markup"] == "pango"
        self.colors = self.get_hex_color_range(self.start_color, self.end_color, 100)

    def run(self):
        if eval(self.display_if):
            if self.lm_sensors_enabled:
                self.output = self.get_output_sensors()
            else:
                self.output = self.get_output_original()

    def get_output_original(self):
        """
        Build the output the original way. Requires no third party libraries.
        """
        with open(self.file, "r") as f:
            temp = float(f.read().strip()) / 1000

        if self.dynamic_color:
            color = self.colors[int(self.percentage(int(temp), self.alert_temp))]
        else:
            color = self.color if temp < self.alert_temp else self.alert_color
        return {
            "full_text": self.format.format(temp=temp),
            "color": color,
        }

    def get_output_sensors(self):
        """
        Build the output using lm_sensors. Requires sensors Python module (see docs).
        """
        data = dict()
        found_sensors = get_sensors()
        if len(found_sensors) == 0:
            raise Exception("No sensors detected! "
                            "Ensure lm-sensors is installed and check the output of the `sensors` command.")
        for sensor in found_sensors:
            data[sensor.name] = self.format_sensor(sensor)
            data["{}_bar".format(sensor.name)] = self.format_sensor_bar(sensor)
        data['temp'] = max((s.current for s in found_sensors))
        return {
            'full_text': self.format.format(**data),
            'urgent': self.get_urgent(found_sensors),
            'color': self.color if not self.dynamic_color else None,
        }

    def get_urgent(self, sensors):
        """ Determine if any sensors should set the urgent flag. """
        if self.urgent_on not in ('warning', 'critical'):
            raise Exception("urgent_on must be one of (warning, critical)")
        for sensor in sensors:
            if self.urgent_on == 'warning' and sensor.is_warning():
                return True
            elif self.urgent_on == 'critical' and sensor.is_critical():
                return True
        return False

    def format_sensor(self, sensor):
        """ Format a sensor value. If pango is enabled color is per sensor. """
        current_val = sensor.current
        if self.pango_enabled:
            percentage = self.percentage(sensor.current, sensor.critical)
            if self.dynamic_color:
                color = self.colors[int(percentage)]
                return self.format_pango(color, current_val)
        return current_val

    def format_sensor_bar(self, sensor):
        """ Build and format a sensor bar. If pango is enabled bar color is per sensor."""
        percentage = self.percentage(sensor.current, sensor.critical)
        bar = make_vertical_bar(int(percentage))
        if self.pango_enabled:
            if self.dynamic_color:
                color = self.colors[int(percentage)]
                return self.format_pango(color, bar)
        return bar

    def format_pango(self, color, value):
        return '{}'.format(color, value)
<commit_msg>temp.py: Fixed bug related to dynamic colors (#672)<commit_after>from i3pystatus import IntervalModule
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_vertical_bar

class Sensor:
    """
    Simple class representing a CPU temperature sensor.
    """

    def __init__(self, name, current, maximum, critical):
        self.name = name.replace(' ', '_')
        self.current = int(current)
        self.maximum = int(maximum) if maximum else int(critical)
        self.critical = int(critical)

    def __repr__(self):
        return "Sensor(name='{}', current={}, maximum={}, critical={})".format(
            self.name,
            self.current,
            self.maximum,
            self.critical
        )

    def is_warning(self):
        return self.current > self.maximum

    def is_critical(self):
        return self.current > self.critical

def get_sensors():
    """ Detect and return a list of Sensor objects """
    import sensors
    found_sensors = list()

    def get_subfeature_value(feature, subfeature_type):
        subfeature = chip.get_subfeature(feature, subfeature_type)
        if subfeature:
            return chip.get_value(subfeature.number)

    for chip in sensors.get_detected_chips():
        for feature in chip.get_features():
            if feature.type == sensors.FEATURE_TEMP:
                try:
                    name = chip.get_label(feature)
                    max = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_MAX)
                    current = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_INPUT)
                    critical = get_subfeature_value(feature, sensors.SUBFEATURE_TEMP_CRIT)
                    if critical:
                        found_sensors.append(Sensor(name=name, current=current, maximum=max, critical=critical))
                except sensors.SensorsException:
                    continue
    return found_sensors

class Temperature(IntervalModule, ColorRangeModule):
    """
    Shows CPU temperature of Intel processors.

    AMD is currently not supported as they can only report a relative temperature, which is pretty useless.

    Requires `colour` module from PyPi

    .. rubric:: Modes of operation

    If lm_sensors_enabled is set to False, the module operates in default mode. This means that:
        * only the {temp} formatter is available
        * alert_temp is honored

    If lm_sensors_enabled is set to True, the module operates in lm_sensors mode. This means that:
        * pysensors must be installed (https://github.com/bastienleonard/pysensors)
        * CPU sensors are discovered dynamically (supporting a sensor per core and multiple CPUs)
        * alert_temp is ignored. The warning or critical values reported by the sensor are used instead (see urgent_on)

    .. rubric:: lm_sensors installation

    In order to take advantage of the lm_sensors library and tools, it must first be installed and configured.

    On Arch this is as simple as:

    .. code-block:: bash

        pacman -S lm_sensors

    On Ubuntu:

    .. code-block:: bash

        sudo apt-get update && sudo apt-get install lm-sensors libsensors4-dev

    The Arch Wiki has a good page on the library - https://wiki.archlinux.org/index.php/lm_sensors

    .. rubric:: lm_sensors_mode formatters

    When ``lm_sensors_enabled`` is True the formatters are created dynamically. In order to discover the formatters that
    are available, it is best to run the sensors command:

    .. code-block:: bash

        ⇒ sensors
        coretemp-isa-0000
        Adapter: ISA adapter
        Physical id 0:  +48.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 0:         +48.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 1:         +46.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 2:         +43.0°C  (high = +80.0°C, crit = +99.0°C)
        Core 3:         +47.0°C  (high = +80.0°C, crit = +99.0°C)

    The module replaces spaces in sensor names with underscores, therefore from the above output we can
    identify the following sensor formatters:

    * Physical_id_0
    * Core_0
    * Core_1
    * Core_2
    * Core_3

    For each sensor a vertical bar is also generated. In this example we would also have the following bars:

    * Physical_id_0_bar
    * Core_0_bar
    * Core_1_bar
    * Core_2_bar
    * Core_3_bar

    Thus, this format string would be valid: "{Physical_id_0}°C {Core_0_bar}{Core_1_bar}{Core_2_bar}{Core_3_bar}"

    .. rubric:: Pango Markup and lm_sensors_mode

    When Pango Markup is enabled and ``dynamic_color`` is True, each sensor's formatter color is displayed independently.
    The color is determined by the proximity of the sensors current value to it's critical value.

    .. rubric:: Example Configuration

    Here is an example configuration based on the sensor values discovered above:

    .. code-block:: python

        status.register("temp",
                        format="{Physical_id_0}°C {Core_0_bar}{Core_1_bar}{Core_2_bar}{Core_3_bar}",
                        hints={"markup": "pango"},
                        lm_sensors_enabled=True,
                        dynamic_color=True)
    """

    settings = (
        ("format",
         "format string used for output. {temp} is the temperature in degrees celsius"),
        ('display_if', 'snippet that gets evaluated. if true, displays the module output'),
        ('lm_sensors_enabled', 'whether or not lm_sensors should be used for obtaining CPU temperature information'),
        ('urgent_on', 'whether to flag as urgent when temperature exceeds urgent value or critical value '
                      '(requires lm_sensors_enabled)'),
        ('dynamic_color', 'whether to set the color dynamically (overrides alert_color)'),
        "color",
        "file",
        "alert_temp",
        "alert_color",
    )
    format = "{temp} °C"
    color = "#FFFFFF"
    file = "/sys/class/thermal/thermal_zone0/temp"
    alert_temp = 90
    alert_color = "#FF0000"
    display_if = 'True'

    lm_sensors_enabled = False
    dynamic_color = False
    urgent_on = 'warning'

    def init(self):
        self.pango_enabled = self.hints.get("markup", False) and self.hints["markup"] == "pango"
        self.colors = self.get_hex_color_range(self.start_color, self.end_color, 100)

    def run(self):
        if eval(self.display_if):
            if self.lm_sensors_enabled:
                self.output = self.get_output_sensors()
            else:
                self.output = self.get_output_original()

    def get_output_original(self):
        """
        Build the output the original way. Requires no third party libraries.
        """
        with open(self.file, "r") as f:
            temp = float(f.read().strip()) / 1000

        if self.dynamic_color:
            perc = int(self.percentage(int(temp), self.alert_temp))
            if (perc > 99):
                perc = 99
            color = self.colors[perc]
        else:
            color = self.color if temp < self.alert_temp else self.alert_color
        return {
            "full_text": self.format.format(temp=temp),
            "color": color,
        }

    def get_output_sensors(self):
        """
        Build the output using lm_sensors. Requires sensors Python module (see docs).
        """
        data = dict()
        found_sensors = get_sensors()
        if len(found_sensors) == 0:
            raise Exception("No sensors detected! "
                            "Ensure lm-sensors is installed and check the output of the `sensors` command.")
        for sensor in found_sensors:
            data[sensor.name] = self.format_sensor(sensor)
            data["{}_bar".format(sensor.name)] = self.format_sensor_bar(sensor)
        data['temp'] = max((s.current for s in found_sensors))
        return {
            'full_text': self.format.format(**data),
            'urgent': self.get_urgent(found_sensors),
            'color': self.color if not self.dynamic_color else None,
        }

    def get_urgent(self, sensors):
        """ Determine if any sensors should set the urgent flag. """
        if self.urgent_on not in ('warning', 'critical'):
            raise Exception("urgent_on must be one of (warning, critical)")
        for sensor in sensors:
            if self.urgent_on == 'warning' and sensor.is_warning():
                return True
            elif self.urgent_on == 'critical' and sensor.is_critical():
                return True
        return False

    def format_sensor(self, sensor):
        """ Format a sensor value. If pango is enabled color is per sensor. """
        current_val = sensor.current
        if self.pango_enabled:
            percentage = self.percentage(sensor.current, sensor.critical)
            if self.dynamic_color:
                color = self.colors[int(percentage)]
                return self.format_pango(color, current_val)
        return current_val

    def format_sensor_bar(self, sensor):
        """ Build and format a sensor bar. If pango is enabled bar color is per sensor."""
        percentage = self.percentage(sensor.current, sensor.critical)
        bar = make_vertical_bar(int(percentage))
        if self.pango_enabled:
            if self.dynamic_color:
                color = self.colors[int(percentage)]
                return self.format_pango(color, bar)
        return bar

    def format_pango(self, color, value):
        return '{}'.format(color, value)
<|endoftext|>