<commit_before>from __future__ import unicode_literals
from django.db import models
from django.test import TestCase
from rest_framework import serializers

class OneToOneTarget(models.Model):
    name = models.CharField(max_length=100)

class OneToOneSource(models.Model):
    name = models.CharField(max_length=100)
    target = models.OneToOneField(OneToOneTarget, related_name='source')

class OneToManyTarget(models.Model):
    name = models.CharField(max_length=100)

class OneToManySource(models.Model):
    name = models.CharField(max_length=100)
    target = models.ForeignKey(OneToManyTarget, related_name='sources')

class ReverseNestedOneToOneTests(TestCase):
    def setUp(self):
        class OneToOneSourceSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToOneSource
                fields = ('id', 'name')

        class OneToOneTargetSerializer(serializers.ModelSerializer):
            source = OneToOneSourceSerializer()

            class Meta:
                model = OneToOneTarget
                fields = ('id', 'name', 'source')

        self.Serializer = OneToOneTargetSerializer

        for idx in range(1, 4):
            target = OneToOneTarget(name='target-%d' % idx)
            target.save()
            source = OneToOneSource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_one_retrieve(self):
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3', 'source': {'id': 3, 'name': 'source-3'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create(self):
        data = {'id': 4, 'name': 'target-4', 'source': {'id': 4, 'name': 'source-4'}}
        serializer = self.Serializer(data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-4')

        # Ensure (target 4, target_source 4, source 4) are added, and
        # everything else is as expected.
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3', 'source': {'id': 3, 'name': 'source-3'}},
            {'id': 4, 'name': 'target-4', 'source': {'id': 4, 'name': 'source-4'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create_with_invalid_data(self):
        data = {'id': 4, 'name': 'target-4', 'source': {'id': 4}}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'source': [{'name': ['This field is required.']}]})

    def test_one_to_one_update(self):
        data = {'id': 3, 'name': 'target-3-updated', 'source': {'id': 3, 'name': 'source-3-updated'}}
        instance = OneToOneTarget.objects.get(pk=3)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-3-updated')

        # Ensure (target 3, target_source 3, source 3) are updated,
        # and everything else is as expected.
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3-updated', 'source': {'id': 3, 'name': 'source-3-updated'}}
        ]
        self.assertEqual(serializer.data, expected)

class ForwardNestedOneToOneTests(TestCase):
    def setUp(self):
        class OneToOneTargetSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToOneTarget
                fields = ('id', 'name')

        class OneToOneSourceSerializer(serializers.ModelSerializer):
            target = OneToOneTargetSerializer()

            class Meta:
                model = OneToOneSource
                fields = ('id', 'name', 'target')

        self.Serializer = OneToOneSourceSerializer

        for idx in range(1, 4):
            target = OneToOneTarget(name='target-%d' % idx)
            target.save()
            source = OneToOneSource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_one_retrieve(self):
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3', 'target': {'id': 3, 'name': 'target-3'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create(self):
        data = {'id': 4, 'name': 'source-4', 'target': {'id': 4, 'name': 'target-4'}}
        serializer = self.Serializer(data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'source-4')

        # Ensure (target 4, target_source 4, source 4) are added, and
        # everything else is as expected.
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3', 'target': {'id': 3, 'name': 'target-3'}},
            {'id': 4, 'name': 'source-4', 'target': {'id': 4, 'name': 'target-4'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create_with_invalid_data(self):
        data = {'id': 4, 'name': 'source-4', 'target': {'id': 4}}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'target': [{'name': ['This field is required.']}]})

    def test_one_to_one_update(self):
        data = {'id': 3, 'name': 'source-3-updated', 'target': {'id': 3, 'name': 'target-3-updated'}}
        instance = OneToOneSource.objects.get(pk=3)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'source-3-updated')

        # Ensure (target 3, target_source 3, source 3) are updated,
        # and everything else is as expected.
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3-updated', 'target': {'id': 3, 'name': 'target-3-updated'}}
        ]
        self.assertEqual(serializer.data, expected)

    # TODO: Nullable 1-1 tests
    # def test_one_to_one_delete(self):
    #     data = {'id': 3, 'name': 'target-3', 'target_source': None}
    #     instance = OneToOneTarget.objects.get(pk=3)
    #     serializer = self.Serializer(instance, data=data)
    #     self.assertTrue(serializer.is_valid())
    #     serializer.save()

    #     # Ensure (target_source 3, source 3) are deleted,
    #     # and everything else is as expected.
    #     queryset = OneToOneTarget.objects.all()
    #     serializer = self.Serializer(queryset)
    #     expected = [
    #         {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
    #         {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
    #         {'id': 3, 'name': 'target-3', 'source': None}
    #     ]
    #     self.assertEqual(serializer.data, expected)

class ReverseNestedOneToManyTests(TestCase):
    def setUp(self):
        class OneToManySourceSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToManySource
                fields = ('id', 'name')

        class OneToManyTargetSerializer(serializers.ModelSerializer):
            sources = OneToManySourceSerializer(many=True, allow_add_remove=True)

            class Meta:
                model = OneToManyTarget
                fields = ('id', 'name', 'sources')

        self.Serializer = OneToManyTargetSerializer

        target = OneToManyTarget(name='target-1')
        target.save()
        for idx in range(1, 4):
            source = OneToManySource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_many_retrieve(self):
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 2, 'name': 'source-2'},
                                                      {'id': 3, 'name': 'source-3'}]},
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_create(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 2, 'name': 'source-2'},
                                                         {'id': 3, 'name': 'source-3'},
                                                         {'id': 4, 'name': 'source-4'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-1')

        # Ensure source 4 is added, and everything else is as
        # expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 2, 'name': 'source-2'},
                                                      {'id': 3, 'name': 'source-3'},
                                                      {'id': 4, 'name': 'source-4'}]}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_create_with_invalid_data(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 2, 'name': 'source-2'},
                                                         {'id': 3, 'name': 'source-3'},
                                                         {'id': 4}]}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'sources': [{}, {}, {}, {'name': ['This field is required.']}]})

    def test_one_to_many_update(self):
        data = {'id': 1, 'name': 'target-1-updated', 'sources': [{'id': 1, 'name': 'source-1-updated'},
                                                                 {'id': 2, 'name': 'source-2'},
                                                                 {'id': 3, 'name': 'source-3'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-1-updated')

        # Ensure (target 1, source 1) are updated,
        # and everything else is as expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1-updated', 'sources': [{'id': 1, 'name': 'source-1-updated'},
                                                              {'id': 2, 'name': 'source-2'},
                                                              {'id': 3, 'name': 'source-3'}]}

        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_delete(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 3, 'name': 'source-3'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        serializer.save()

        # Ensure source 2 is deleted, and everything else is as
        # expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 3, 'name': 'source-3'}]}

        ]
        self.assertEqual(serializer.data, expected)
<commit_msg>Fix deprecation warnings in relations_nested tests<commit_after>from __future__ import unicode_literals
from django.db import models
from django.test import TestCase
from rest_framework import serializers

class OneToOneTarget(models.Model):
    name = models.CharField(max_length=100)

class OneToOneSource(models.Model):
    name = models.CharField(max_length=100)
    target = models.OneToOneField(OneToOneTarget, related_name='source')

class OneToManyTarget(models.Model):
    name = models.CharField(max_length=100)

class OneToManySource(models.Model):
    name = models.CharField(max_length=100)
    target = models.ForeignKey(OneToManyTarget, related_name='sources')

class ReverseNestedOneToOneTests(TestCase):
    def setUp(self):
        class OneToOneSourceSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToOneSource
                fields = ('id', 'name')

        class OneToOneTargetSerializer(serializers.ModelSerializer):
            source = OneToOneSourceSerializer()

            class Meta:
                model = OneToOneTarget
                fields = ('id', 'name', 'source')

        self.Serializer = OneToOneTargetSerializer

        for idx in range(1, 4):
            target = OneToOneTarget(name='target-%d' % idx)
            target.save()
            source = OneToOneSource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_one_retrieve(self):
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3', 'source': {'id': 3, 'name': 'source-3'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create(self):
        data = {'id': 4, 'name': 'target-4', 'source': {'id': 4, 'name': 'source-4'}}
        serializer = self.Serializer(data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-4')

        # Ensure (target 4, target_source 4, source 4) are added, and
        # everything else is as expected.
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3', 'source': {'id': 3, 'name': 'source-3'}},
            {'id': 4, 'name': 'target-4', 'source': {'id': 4, 'name': 'source-4'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create_with_invalid_data(self):
        data = {'id': 4, 'name': 'target-4', 'source': {'id': 4}}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'source': [{'name': ['This field is required.']}]})

    def test_one_to_one_update(self):
        data = {'id': 3, 'name': 'target-3-updated', 'source': {'id': 3, 'name': 'source-3-updated'}}
        instance = OneToOneTarget.objects.get(pk=3)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-3-updated')

        # Ensure (target 3, target_source 3, source 3) are updated,
        # and everything else is as expected.
        queryset = OneToOneTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
            {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
            {'id': 3, 'name': 'target-3-updated', 'source': {'id': 3, 'name': 'source-3-updated'}}
        ]
        self.assertEqual(serializer.data, expected)

class ForwardNestedOneToOneTests(TestCase):
    def setUp(self):
        class OneToOneTargetSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToOneTarget
                fields = ('id', 'name')

        class OneToOneSourceSerializer(serializers.ModelSerializer):
            target = OneToOneTargetSerializer()

            class Meta:
                model = OneToOneSource
                fields = ('id', 'name', 'target')

        self.Serializer = OneToOneSourceSerializer

        for idx in range(1, 4):
            target = OneToOneTarget(name='target-%d' % idx)
            target.save()
            source = OneToOneSource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_one_retrieve(self):
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3', 'target': {'id': 3, 'name': 'target-3'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create(self):
        data = {'id': 4, 'name': 'source-4', 'target': {'id': 4, 'name': 'target-4'}}
        serializer = self.Serializer(data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'source-4')

        # Ensure (target 4, target_source 4, source 4) are added, and
        # everything else is as expected.
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3', 'target': {'id': 3, 'name': 'target-3'}},
            {'id': 4, 'name': 'source-4', 'target': {'id': 4, 'name': 'target-4'}}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_one_create_with_invalid_data(self):
        data = {'id': 4, 'name': 'source-4', 'target': {'id': 4}}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'target': [{'name': ['This field is required.']}]})

    def test_one_to_one_update(self):
        data = {'id': 3, 'name': 'source-3-updated', 'target': {'id': 3, 'name': 'target-3-updated'}}
        instance = OneToOneSource.objects.get(pk=3)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'source-3-updated')

        # Ensure (target 3, target_source 3, source 3) are updated,
        # and everything else is as expected.
        queryset = OneToOneSource.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'source-1', 'target': {'id': 1, 'name': 'target-1'}},
            {'id': 2, 'name': 'source-2', 'target': {'id': 2, 'name': 'target-2'}},
            {'id': 3, 'name': 'source-3-updated', 'target': {'id': 3, 'name': 'target-3-updated'}}
        ]
        self.assertEqual(serializer.data, expected)

    # TODO: Nullable 1-1 tests
    # def test_one_to_one_delete(self):
    #     data = {'id': 3, 'name': 'target-3', 'target_source': None}
    #     instance = OneToOneTarget.objects.get(pk=3)
    #     serializer = self.Serializer(instance, data=data)
    #     self.assertTrue(serializer.is_valid())
    #     serializer.save()

    #     # Ensure (target_source 3, source 3) are deleted,
    #     # and everything else is as expected.
    #     queryset = OneToOneTarget.objects.all()
    #     serializer = self.Serializer(queryset)
    #     expected = [
    #         {'id': 1, 'name': 'target-1', 'source': {'id': 1, 'name': 'source-1'}},
    #         {'id': 2, 'name': 'target-2', 'source': {'id': 2, 'name': 'source-2'}},
    #         {'id': 3, 'name': 'target-3', 'source': None}
    #     ]
    #     self.assertEqual(serializer.data, expected)

class ReverseNestedOneToManyTests(TestCase):
    def setUp(self):
        class OneToManySourceSerializer(serializers.ModelSerializer):
            class Meta:
                model = OneToManySource
                fields = ('id', 'name')

        class OneToManyTargetSerializer(serializers.ModelSerializer):
            sources = OneToManySourceSerializer(many=True, allow_add_remove=True)

            class Meta:
                model = OneToManyTarget
                fields = ('id', 'name', 'sources')

        self.Serializer = OneToManyTargetSerializer

        target = OneToManyTarget(name='target-1')
        target.save()
        for idx in range(1, 4):
            source = OneToManySource(name='source-%d' % idx, target=target)
            source.save()

    def test_one_to_many_retrieve(self):
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 2, 'name': 'source-2'},
                                                      {'id': 3, 'name': 'source-3'}]},
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_create(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 2, 'name': 'source-2'},
                                                         {'id': 3, 'name': 'source-3'},
                                                         {'id': 4, 'name': 'source-4'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-1')

        # Ensure source 4 is added, and everything else is as
        # expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 2, 'name': 'source-2'},
                                                      {'id': 3, 'name': 'source-3'},
                                                      {'id': 4, 'name': 'source-4'}]}
        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_create_with_invalid_data(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 2, 'name': 'source-2'},
                                                         {'id': 3, 'name': 'source-3'},
                                                         {'id': 4}]}
        serializer = self.Serializer(data=data)
        self.assertFalse(serializer.is_valid())
        self.assertEqual(serializer.errors, {'sources': [{}, {}, {}, {'name': ['This field is required.']}]})

    def test_one_to_many_update(self):
        data = {'id': 1, 'name': 'target-1-updated', 'sources': [{'id': 1, 'name': 'source-1-updated'},
                                                                 {'id': 2, 'name': 'source-2'},
                                                                 {'id': 3, 'name': 'source-3'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        obj = serializer.save()
        self.assertEqual(serializer.data, data)
        self.assertEqual(obj.name, 'target-1-updated')

        # Ensure (target 1, source 1) are updated,
        # and everything else is as expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1-updated', 'sources': [{'id': 1, 'name': 'source-1-updated'},
                                                              {'id': 2, 'name': 'source-2'},
                                                              {'id': 3, 'name': 'source-3'}]}

        ]
        self.assertEqual(serializer.data, expected)

    def test_one_to_many_delete(self):
        data = {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                         {'id': 3, 'name': 'source-3'}]}
        instance = OneToManyTarget.objects.get(pk=1)
        serializer = self.Serializer(instance, data=data)
        self.assertTrue(serializer.is_valid())
        serializer.save()

        # Ensure source 2 is deleted, and everything else is as
        # expected.
        queryset = OneToManyTarget.objects.all()
        serializer = self.Serializer(queryset, many=True)
        expected = [
            {'id': 1, 'name': 'target-1', 'sources': [{'id': 1, 'name': 'source-1'},
                                                      {'id': 3, 'name': 'source-3'}]}

        ]
        self.assertEqual(serializer.data, expected)
<|endoftext|>


<commit_before>"""
``revscoring extract_features -h``
::

    Adds features to a set of labeled revisions.

    Reads a TSV file of <rev_id>\t<label> pairs and replaces the
    <rev_id> field with the extracted feature values.

    Input: <rev_id>[TAB]<label>

    Output: <feature_1>[TAB]<feature_2>[TAB]...[TAB]<label>

    Usage:
        extract_features -h | --help
        extract_features <features> --host=<url> [--rev-labels=<path>]
                                                 [--value-labels=<path>]
                                                 [--include-revid]
                                                 [--verbose] [--debug]

    Options:
        -h --help                Print this documentation
        <features>               Classpath to a list/tuple of features
        --host=<url>             The url pointing to a MediaWiki API to use
                                 for extracting features
        --rev-labels=<path>      Path to a file containing rev_id-label pairs
                                 [default: <stdin>]
        --value-labels=<path>    Path to a file to write feature-labels to
                                 [default: <stdout>]
        --include-revid          If set, include the revision ID as the first
                                 column in the output TSV
        --verbose                Print dots and stuff
        --debug                  Print debug logging
"""
import logging
import sys
from itertools import islice

import docopt
import mwapi

from ..errors import RevisionNotFound
from ..extractors import APIExtractor
from .util import encode, import_from_path

logger = logging.getLogger(__name__)

BATCH_SIZE = 50

def main(argv=None):
    args = docopt.docopt(__doc__, argv=argv)

    logging.basicConfig(
        level=logging.WARNING if not args['--debug'] else logging.DEBUG,
        format='%(asctime)s %(levelname)s:%(name)s -- %(message)s'
    )

    features = import_from_path(args['<features>'])

    session = mwapi.Session(args['--host'],
                            user_agent="Revscoring feature extractor utility")
    extractor = APIExtractor(session)

    if args['--rev-labels'] == "<stdin>":
        rev_labels = read_rev_labels(sys.stdin)
    else:
        rev_labels = read_rev_labels(open(args['--rev-labels']))

    if args['--value-labels'] == "<stdout>":
        value_labels = sys.stdout
    else:
        value_labels = open(args['--value-labels'], 'w')

    include_revid = bool(args['--include-revid'])

    verbose = args['--verbose']
    debug = args['--debug']

    run(rev_labels, value_labels, features, extractor, include_revid,
        verbose, debug)

def read_rev_labels(f):
    # Check if first line is a header
    rev_id, label = f.readline().strip().split("\t")
    if rev_id != "rev_id":
        yield int(rev_id), label

    for line in f:
        rev_id, label = line.strip().split('\t')
        yield int(rev_id), label

def run(rev_labels, value_labels, features, extractor, include_revid,
        verbose, debug):
    logging.basicConfig(
        level=logging.WARNING if not debug else logging.DEBUG,
        format='%(asctime)s %(levelname)s:%(name)s -- %(message)s'
    )

    while True:
        batch_rev_labels = list(islice(rev_labels, BATCH_SIZE))
        if len(batch_rev_labels) == 0:
            break

        rev_ids, labels = zip(*batch_rev_labels)

        error_values_label = zip(extractor.extract(rev_ids, features),
                                 batch_rev_labels)
        for (error, values), (rev_id, label) in error_values_label:
            try:
                if isinstance(error, RevisionNotFound):
                    if verbose:
                        sys.stderr.write("?")
                        sys.stderr.flush()
                elif error is not None:
                    logger.error("An error occured while processing {0}:"
                                 .format(rev_id))
                    logger.error("\t{0}: {1}"
                                 .format(error.__class__.__name__,
                                         str(error)))
                else:
                    fields = list(values) + [label]

                    if include_revid:
                        fields = [rev_id] + fields

                    value_labels.write("\t".join(encode(v)
                                                 for v in fields))
                    value_labels.write("\n")

                    if verbose:
                        sys.stderr.write(".")
                        sys.stderr.flush()

            except KeyboardInterrupt:
                sys.stderr.write("^C detected.  Shutting down.\n")
                break

    if verbose:
        sys.stderr.write("\n")
<commit_msg>Adds parallelism to feature extraction<commit_after>"""
``revscoring extract_features -h``
::

    Adds features to a set of labeled revisions.

    Reads a TSV file of <rev_id>\t<label> pairs and replaces the
    <rev_id> field with the extracted feature values.

    Input: <rev_id>[TAB]<label>

    Output: <feature_1>[TAB]<feature_2>[TAB]...[TAB]<label>

    Usage:
        extract_features -h | --help
        extract_features <features> --host=<url> [--rev-labels=<path>]
                                                 [--value-labels=<path>]
                                                 [--include-revid]
                                                 [--extractors=<num>]
                                                 [--verbose] [--debug]

    Options:
        -h --help               Print this documentation
        <features>              Classpath to a list/tuple of features
        --host=<url>            The url pointing to a MediaWiki API to use
                                for extracting features
        --rev-labels=<path>     Path to a file containing rev_id-label pairs
                                [default: <stdin>]
        --value-labels=<path>   Path to a file to write feature-labels to
                                [default: <stdout>]
        --include-revid         If set, include the revision ID as the first
                                column in the output TSV
        --extractors=<num>      The number of extractors to run in parallel
                                [default: <cpu count>]
        --verbose               Print dots and stuff
        --debug                 Print debug logging
"""
import logging
import sys
from itertools import islice
from multiprocessing import Pool, cpu_count

import docopt
import mwapi

from ..errors import RevisionNotFound
from ..extractors import APIExtractor
from .util import encode, import_from_path

logger = logging.getLogger(__name__)

def main(argv=None):
    args = docopt.docopt(__doc__, argv=argv)

    logging.basicConfig(
        level=logging.WARNING if not args['--debug'] else logging.DEBUG,
        format='%(asctime)s %(levelname)s:%(name)s -- %(message)s'
    )

    features = import_from_path(args['<features>'])

    session = mwapi.Session(args['--host'],
                            user_agent="Revscoring feature extractor utility")
    extractor = APIExtractor(session)

    if args['--rev-labels'] == "<stdin>":
        rev_labels = read_rev_labels(sys.stdin)
    else:
        rev_labels = read_rev_labels(open(args['--rev-labels']))

    if args['--value-labels'] == "<stdout>":
        value_labels = sys.stdout
    else:
        value_labels = open(args['--value-labels'], 'w')

    include_revid = bool(args['--include-revid'])

    if args['--extractors'] == "<cpu count>":
        extractors = cpu_count()
    else:
        extractors = int(extractors)

    verbose = args['--verbose']
    debug = args['--debug']

    run(rev_labels, value_labels, features, extractor, include_revid,
        extractors, verbose, debug)

def read_rev_labels(f):
    # Check if first line is a header
    rev_id, label = f.readline().strip().split("\t")
    if rev_id != "rev_id":
        yield int(rev_id), label

    for line in f:
        rev_id, label = line.strip().split('\t')
        yield int(rev_id), label

def run(rev_labels, value_labels, features, extractor, include_revid,
        extractors, verbose, debug):
    logging.basicConfig(
        level=logging.WARNING if not debug else logging.DEBUG,
        format='%(asctime)s %(levelname)s:%(name)s -- %(message)s'
    )

    extractor_context = ConfiguredExtractor(extractor, features)
    extractor_pool = Pool(processes=extractors)

    error_rev_label_features = extractor_pool.imap(extractor_context.extract,
                                                   rev_labels)

    for error, rev_id, label, feature_values in error_rev_label_features:
        if isinstance(error, RevisionNotFound):
            if verbose:
                sys.stderr.write("?")
                sys.stderr.flush()
        elif error is not None:
            logger.error("An error occured while processing {0}:"
                         .format(rev_id))
            logger.error("\t{0}: {1}"
                         .format(error.__class__.__name__,
                                 str(error)))
        else:
            fields = list(feature_values) + [label]

            if include_revid:
                fields = [rev_id] + fields

            value_labels.write("\t".join(encode(v)
                                         for v in fields))
            value_labels.write("\n")

            if verbose:
                sys.stderr.write(".")
                sys.stderr.flush()

    if verbose:
        sys.stderr.write("\n")

class ConfiguredExtractor:

    def __init__(self, extractor, features):
        self.extractor = extractor
        self.features = features

    def extract(self, rev_label):
        rev_id, label = rev_label
        try:
            feature_values = list(self.extractor.extract(rev_id, self.features))
            error = None
        except Exception as e:
            feature_values = None
            error = e

        return error, rev_id, label, feature_values
<|endoftext|>


<commit_before>try:
    import wpilib
except ImportError:
    from pyfrc import wpilib
from common.autonomous_helper import StatefulAutonomous, timed_state

class TwoBall(StatefulAutonomous):

    MODE_NAME = 'Two Balls Hot'
    DEFAULT = False

    def __init__(self, components):
        super().__init__(components)

        '''#self.register_sd_var('SecondBallDriveTime',1.5);
        self.register_sd_var('SecondBallShoot',1.5);
        #self.register_sd_var('GetSecondBallTime',0.7);
        self.register_sd_var('DriveRotateTime2',0.1);
        self.register_sd_var('DriveRotateTime3',0.1);'''
        self.register_sd_var('hotLeft',False)
        self.register_sd_var("hotRight",False)

        wpilib.SmartDashboard.PutNumber('IsHotLeft', False)
        wpilib.SmartDashboard.PutNumber('IsHotRight', False)
        self.decided = False
    def on_enable(self):
        '''This function is called when autonomous mode is enabled'''
        super().on_enable()
        '''#self.drive_time_getSecondBall = wpilib.SmartDashboard.GetNumber("SecondBallDriveTime")
        #self.drive_time_reverse = wpilib.SmartDashboard.GetNumber("GetSecondBallTime")
        #self.drive_time_shoot2 = wpilib.SmartDashboard.GetNumber("SecondBallShoot")
        self.drive_rotate_time2 = wpilib.SmartDashboard.GetNumber('DriveRotateTime2')
        self.drive_rotate_time3 = wpilib.SmartDashboard.GetNumber('DriveRotateTime3')'''

    def update(self, tm):
        if tm > 0.3:
            self.catapult.pulldown()
            print(tm, update)
        if not self.decided:
            self.hotLeft = wpilib.SmartDashboard.GetBoolean("IsHotLeft")
            self.hotRight = wpilib.SmartDashboard.GetBoolean("IsHotRight")

            if (self.hotLeft or self.hotRight) and not (self.hotLeft and self.hotRight):
                self.decided = True

                if self.hotLeft:
                    next_state("RotateLeft")
                    #print ("hot Left")
                else:
                    next_state("RotateRight")
                    #print ('hot right')
            else:
                next_state("RotateLeft")
        super().update(tm)

    def RotateRight(self, tm, state_tm):
        self.drive_rotate_speed = self.drive_rotate_speed_right
        print('hot left')
        print("RotateRight", tm)
    def RotateLeft(self,tm, state_tm):
        self.drive_rotate_speed = self.drive_rotate_speed_left
        print('hot right')
        print("RotateLeft", tm)

    #
    #I have no idea how the timing is supposed to work from states 'nextball1' to 'launch2'.
    #For the moment I'm just going to add one second to each state.
    #
    #Please contact Timmy and get him to explain the timngs on his old TwoBall class
    @timed_state(duration=.5, next_state='drive_wait', first=True)
    def drive_wait(self, tm, state_tm):
        self.intake.armDown
        print("drive_wait'", tm)

    @timed_state(duration=1,next_state='drive_rotate')
    def drive_rotate(self, tm, state_tm):
        self.drive.move(0,0,self.drive_rotate_speed)
        print("drive_rotate", tm)
    @timed_state(duration=1,next_state='drive_start')
    def drive_start(self, tm, state_tm):
         self.drive.move(0, self.drive_speed, 0)
         print("screwleondrive_start", tm)
    @timed_state(duration=1,next_state='launch')
    def launch(self, tm, state_tm):
        self.catapult.launchNoSensor()
        print("launch", tm)

    @timed_state(duration=1, next_state='next_ball1')
    def next_ball1(self,tm, state_tm):
            self.drive.move(0, -1*self.drive_speed, 0)
            self.intake.ballIn()
            print("next_ball1", tm)
    @timed_state(duration=1, next_state='next_ball1_rotate')
    def next_ball1_rotate(self,tm, state_tm):
            self.drive.move(0, 0, -1*self.drive_rotate_speed)
            self.intake.ballIn()
            print("next_ball1_rotate", tm)
    @timed_state(duration=1, next_state='next_ball2')
    def next_ball2(self,tm, state_tm):
            self.drive.move(0, self.drive_speed, 0)
            self.intake.ballIn()
            print("next_ball2", tm)
    @timed_state(duration=1,next_state='rotate2')
    def rotate2(self,tm, state_tm):
        self.drive.move(0,self.drive_speed,0)
        self.intake.ballIn()

    @timed_state(duration=1, next_state='launch2')
    def launch2(self,tm, state_tm):
            # Finally, fire and keep firing for 1 seconds
            self.catapult.launchNoSensor()
    <commit_msg>TwoBallHot passes tests<commit_after>try:
    import wpilib
except ImportError:
    from pyfrc import wpilib
from common.autonomous_helper import StatefulAutonomous, timed_state

class TwoBall(StatefulAutonomous):

    MODE_NAME = 'Two Balls Hot'
    DEFAULT = False

    def __init__(self, components):
        super().__init__(components)

        '''#self.register_sd_var('SecondBallDriveTime',1.5);
        self.register_sd_var('SecondBallShoot',1.5);
        #self.register_sd_var('GetSecondBallTime',0.7);
        self.register_sd_var('DriveRotateTime2',0.1);
        self.register_sd_var('DriveRotateTime3',0.1);'''
        self.register_sd_var('hotLeft',False)
        self.register_sd_var("hotRight",False)
        #please change the drive rotate speeds, I'm really
        #just guessing about the values
        self.register_sd_var('drive_rotate_speed_left',1)
        self.register_sd_var('drive_rotate_speed_right',-1)

        wpilib.SmartDashboard.PutNumber('IsHotLeft', False)
        wpilib.SmartDashboard.PutNumber('IsHotRight', False)
        self.decided = False
    #def on_enable(self):
        '''This function is called when autonomous mode is enabled'''
        #super().on_enable()

        '''#self.drive_time_getSecondBall = wpilib.SmartDashboard.GetNumber("SecondBallDriveTime")
        #self.drive_time_reverse = wpilib.SmartDashboard.GetNumber("GetSecondBallTime")
        #self.drive_time_shoot2 = wpilib.SmartDashboard.GetNumber("SecondBallShoot")
        self.drive_rotate_time2 = wpilib.SmartDashboard.GetNumber('DriveRotateTime2')
        self.drive_rotate_time3 = wpilib.SmartDashboard.GetNumber('DriveRotateTime3')'''

    def update(self, tm):
        if tm > 0.3:
            self.catapult.pulldown()

        if not self.decided:
            self.hotLeft = wpilib.SmartDashboard.GetBoolean("IsHotLeft")
            self.hotRight = wpilib.SmartDashboard.GetBoolean("IsHotRight")

            if (self.hotLeft or self.hotRight) and not (self.hotLeft and self.hotRight):
                self.decided = True

                if self.hotLeft:
                    self.next_state("RotateLeft")
                    #print ("hot Left")
                else:
                    self.next_state("RotateRight")
                    #print ('hot right')
            else:
                self.next_state("RotateLeft")
        super().update(tm)
    @timed_state(duration=1)
    def RotateRight(self, tm, state_tm):
        self.drive_rotate_speed = self.drive_rotate_speed_right
        print('hot left')
        print("RotateRight", tm)
    @timed_state(duration=1)
    def RotateLeft(self,tm, state_tm):
        self.drive_rotate_speed = self.drive_rotate_speed_left
        print('hot right')
        print("RotateLeft", tm)

    #
    #I have no idea how the timing is supposed to work from states 'nextball1' to 'launch2'.
    #For the moment I'm just going to add one second to each state.
    #
    #Please contact Timmy and get him to explain the timngs on his old TwoBall class
    @timed_state(duration=.5, next_state='drive_wait', first=True)
    def drive_wait(self, tm, state_tm):
        self.intake.armDown
        print("drive_wait'", tm)

    @timed_state(duration=1,next_state='drive_rotate')
    def drive_rotate(self, tm, state_tm):
        self.drive.move(0,0,self.drive_rotate_speed)
        print("drive_rotate", tm)
    @timed_state(duration=1,next_state='drive_start')
    def drive_start(self, tm, state_tm):
         self.drive.move(0, self.drive_speed, 0)
         print("screwleondrive_start", tm)
    @timed_state(duration=1,next_state='launch')
    def launch(self, tm, state_tm):
        self.catapult.launchNoSensor()
        print("launch", tm)

    @timed_state(duration=1, next_state='next_ball1')
    def next_ball1(self,tm, state_tm):
            self.drive.move(0, -1*self.drive_speed, 0)
            self.intake.ballIn()
            print("next_ball1", tm)
    @timed_state(duration=1, next_state='next_ball1_rotate')
    def next_ball1_rotate(self,tm, state_tm):
            self.drive.move(0, 0, -1*self.drive_rotate_speed)
            self.intake.ballIn()
            print("next_ball1_rotate", tm)
    @timed_state(duration=1, next_state='next_ball2')
    def next_ball2(self,tm, state_tm):
            self.drive.move(0, self.drive_speed, 0)
            self.intake.ballIn()
            print("next_ball2", tm)
    @timed_state(duration=1,next_state='rotate2')
    def rotate2(self,tm, state_tm):
        self.drive.move(0,self.drive_speed,0)
        self.intake.ballIn()

    @timed_state(duration=1, next_state='launch2')
    def launch2(self,tm, state_tm):
            # Finally, fire and keep firing for 1 seconds
            self.catapult.launchNoSensor()
    <|endoftext|>


<commit_before>"""
Author: NekoGamiYuki
Version: 1.0.0

Description:
A way to create and manage deaths during a game. Specifically made for quick
managing of deaths during livestreams.
"""

# Imported Modules--------------------------------------------------------------
import re
import sqlite3
import warnings
import os
from .. import IRC
from . import module_tools

# TODO: Allow for switching games.
# TODO: Allow decrementing and incrementing by specific values
#       Mods could do !dg for death_game and then name it.

# Global Variables--------------------------------------------------------------
default_config = {
    'current_game': '',
    'nickname': ''
}

default_data = {
    'deaths': {}
}

# Regex-------------------------------------------------------------------------
deathcount_regex = re.compile(r"dc( \w+)?")
increment_regex = re.compile(r"d( \d+)?")
decrement_regex = re.compile(r"dd( \d+)?")
add_game_regex = re.compile(r"deathcount add game (\w+|\W+)")
remove_game_regex = re.compile(r"deathcount remove game (\w+|\W+)")
set_game_regex = re.compile(r"deathcount set game (\w+|\W+)")
set_nickname_regex = re.compile(r"deathcount set nickname (\w+)")

# Module Registration-----------------------------------------------------------
module_tools.register_module()
logger = module_tools.create_logger()
storage_directory = module_tools.get_storage_directory()

# Database----------------------------------------------------------------------
connection = sqlite3.connect(os.path.join(storage_directory, __name__))
cursor = connection.cursor()

# Setup database on startup
# Setting up nickname table
# NOTE: This is not affected by channels as it is capitalized
cursor.execute(
    "CREATE TABLE IF NOT EXISTS nicknames "
    "(nickname TEXT, channel TEXT)"
)

cursor.execute(
    "CREATE TABLE IF NOT EXISTS deathcount "
    "(game TEXT, count INT, isdefault TEXT, channel TEXT)"
)

# Nickname management-----------------------------------------------------------
def db_set_nickname(nickname: str, channel: str):
    cursor.execute(
        "SELECT * FROM Nicknames "
        "WHERE channel=(?)",
        (channel,)
    )
    data = cursor.fetchall()
    if not data:
        cursor.execute(
            "INSERT INTO Nicknames "
            "(nickname, channel) "
            "VALUES (?, ?)",
            (nickname, channel)
        )
    else:
        cursor.execute(
            "UPDATE Nicknames "
            "SET nickname=(?) "
            "WHERE channel=(?)",
            (nickname, channel.lower())
        )

    connection.commit()

def db_get_streamer_nickname(channel: str) -> str:
    """Returns channel nickname, if not found it returns the channel's name
    """
    cursor.execute(
        "SELECT nickname FROM Nicknames "
        "WHERE channel=?",
        (channel.lower(),)
    )
    data = cursor.fetchall()
    if not data:
        nickname = ''
    else:
        nickname = data[0][0]

    return nickname

# Data management---------------------------------------------------------------
def db_check_for_duplicate(game: str, channel: str) -> bool:
    """Checks if a game is already in the database for a channel
    """
    cursor.execute(
        "SELECT * FROM deathcount "
        "WHERE channel=(?)AND game=(?)",
        (channel.lower(), game.lower())
    )
    data = cursor.fetchall()
    if len(data) > 0:
        return False
    else:
        return True

def db_check_game_exists(game: str, channel: str) -> bool:
    """Wrapper for db_check_for_duplicates, for readability

    Reverses the return of db_check_for_duplicates, True indicates
    that a game exists, False indicates that it does not.
    """
    return not db_check_for_duplicate(game, channel)

def db_get_death_count(game: str, channel: str) -> int:
    """Returns the deathcount for a game in a specific channel.
    """
    cursor.execute(
        "SELECT count FROM deathcount "
        "WHERE channel=(?) AND game=(?)",
        (channel.lower(), game.lower())
    )
    data = cursor.fetchall()
    if not data:
        return -1
    else:
        return data[0][0]

def db_get_game_count(channel: str):
    """Returns the number of games in a channel
    """
    cursor.execute(
        "SELECT game FROM deathcount "
        "WHERE channel=(?)",
        (channel.lower(),)
    )
    data = cursor.fetchall()
    return len(data)

def db_add_game(game: str, channel: str, set_default_game: bool=False) -> bool:
    """Adds a game to the database for a channel, does not allow duplicates
    """
    if not db_check_for_duplicate(game, channel):
        return False
    else:
        if set_default_game:
            default = "YES"
        else:
            default = "NO"

        cursor.execute(
            "INSERT INTO deathcount "
            "(game, count, isdefault, channel) "
            "VALUES (?, ?, ?, ?)",
            (game.lower(), 0, default, channel.lower())
        )
        connection.commit()

        return True

def db_remove_game(game: str, channel: str) -> bool:
    """Removes a game from the database, for a specific channel
    """
    if db_check_game_exists(game, channel):
        cursor.execute(
            "DELETE FROM deathcount "
            "WHERE channel=(?) AND game=(?)",
            (channel.lower(), game.lower())
        )
        connection.commit()
        return True
    else:
        return False

def db_increment_count(game: str, channel: str, amount: int) -> bool:
    """Incremenets the deathcount of a game, for a specific channel
    """

    if db_check_game_exists(game, channel):
        cursor.execute(
            "UPDATE deathcount "
            "SET count=count+(?) "
            "WHERE channel=(?) AND game=(?)",
            (amount, channel.lower(), game.lower())
        )
        connection.commit()
        return True
    else:
        return False

def db_decrement_count(game: str, channel: str, amount: int) -> bool:
    if db_check_game_exists(game, channel):
        count = db_get_death_count(game, channel)

        if count - amount < 0:
            cursor.execute(
                "UPDATE deathcount "
                "SET count=(?) "
                "WHERE channel=(?) AND game=(?)",
                (0, channel.lower(), game.lower())
            )
            connection.commit()
            return True
        else:
            cursor.execute(
                "UPDATE deathcount "
                "SET count=count-(?) "
                "WHERE channel=(?) AND game=(?)",
                (amount, channel.lower(), game.lower())
            )
            connection.commit()
            return True
    else:
        return False

def db_reset_count(game: str, channel: str) -> bool:
    """Sets the counter to 0 for a game, in a specific channel
    """
    if db_check_game_exists(game, channel):
        cursor.execute(
            "UPDATE deathcount "
            "SET count=(?) "
            "WHERE channel=(?) AND game=(?)",
            (0, channel.lower(), game.lower())
        )
        connection.commit()
        return True
    else:
        return False

def db_get_default_game(channel: str):
    """Returns the default game for a specific channel

    The default game is the game that will be used whenever a command
    is called without any game specified.
    """
    cursor.execute(
        "SELECT game FROM deathcount "
        "WHERE channel=(?) AND isdefault=(?)",
        (channel.lower(), "YES")
    )
    data = cursor.fetchall()

    if len(data) == 0:
        return ''
    else:
        return data[0][0]

def db_set_default_game(game: str, channel:str) -> bool:
    """Marks a game as being the default for a specfic channel
    """

    if db_check_game_exists(game, channel):
        # Unset previous default
        cursor.execute(
            "UPDATE deathcount "
            "SET isdefault='NO' "
            "WHERE channel=(?) AND isdefault=(?)",
            (channel.lower(), "YES")
        )
        connection.commit()

        # Set new game to default
        cursor.execute(
            "UPDATE deathcount "
            "SET isdefault=(?) "
            "WHERE channel=(?) AND game=(?)",
            ("YES", channel.lower(), game.lower())
        )
        connection.commit()
        return True
    else:
        return False

def db_close_connection():
    logger.info(
        "Saving database and closing connection."
    )
    cursor.close()
    connection.close()

# Command functions ------------------------------------------------------------
def increment_deaths(user: IRC.User):

    user_input = user.message.split(module_tools.DEFAULT_COMMAND_PREFIX, 1)[1]
    parsed_input = increment_regex.findall(user_input)

    if parsed_input[0]:
        try:
            increment_by = int(parsed_input[0].strip())
        except ValueError:
            # I doubt this will ever launch, but I'd prefer not to have the bot
            # crash.
            logger.warning(
                "Failed to convert user's input into an integer! "
                "User input: {}".format(user.message)
            )
            return
    else:
        increment_by = 1

    game = db_get_default_game(user.chatted_from)

    if not game:
        user.send_message(
            "No game has been set for the death counter. The streamer or a mod"
            "can set the game by using '!deathcount set game <game>'"
        )

    if not db_increment_count(game, user.chatted_from, increment_by):
        user.send_message(
            "The game you've given is not in the list of games!"
        )
    else:
        logger.info(
            "{} has Incremented death count by {}, for the game {}, in the channel "
            "{}.".format(user.name, increment_by, game, user.chatted_from)
        )

        user.send_message(
            "Death count has been incremented by {}, for the "
            "game '{}'. RIP sfhSAD".format(increment_by, game)
        )

def decrement_deaths(user: IRC.User):
    user_input = user.message.split(module_tools.DEFAULT_COMMAND_PREFIX, 1)[1]
    parsed_input = decrement_regex.findall(user_input)

    if parsed_input[0]:
        try:
            decrement_by = int(parsed_input[0].strip())
        except ValueError:
            # I doubt this will ever launch, but I'd prefer not to have the bot
            # crash.
            logger.warning(
                "Failed to convert user's input into an integer! "
                "User input: {}".format(user.message)
            )
            return
    else:
        decrement_by = 1

    game = db_get_default_game(user.chatted_from)

    if not game:
        user.send_message(
            "No game has been set for the death counter. The streamer or a mod"
            "can set the game by using '!deathcount set game <game>'"
        )

    if not db_decrement_count(game, user.chatted_from, decrement_by):
        user.send_message(
            "The game you've given is not in the list of games!"
        )
    else:
        logger.info(
            "{} has decremented the death count for the game '{}' by {} in the channel "
            "'{}'.".format(user.name, game, decrement_by, user.chatted_from)
        )

        user.send_message(
            "Death count has been decremented by {} for the game '{}'. "
            "sfhOH".format(decrement_by, game)
        )

def reset_deaths(user: IRC.User):
    game = db_get_default_game(user.chatted_from)

    if not game:
        user.send_message(
            "No game has been set for the death counter. The streamer or a mod "
            "can set the game by using '!deathcount set game <game>'"
        )

    if not db_reset_count(game, user.chatted_from):
        user.send_message(
            "The game currently set is not in the database. This should not "
            "have happened... OhGod"
        )
    else:
        logger.info("Reset death count for game '{}' in channel {}.".format(
            game, user.chatted_from)
        )

        user.send_message(
            "Deaths have been reset to zero for the game '{}'. "
            "sfhWOW".format(game)
        )

def death_count(user: IRC.User):
    user_input = user.message.split(module_tools.DEFAULT_COMMAND_PREFIX, 1)[1]
    parsed_input = deathcount_regex.findall(user_input)

    if parsed_input[0]:
        game = parsed_input[0].strip().lower()
    else:
        game = db_get_default_game(user.chatted_from)

    if not game:
        user.send_message(
            "No game has been set for the death counter. The streamer or a mod "
            "can set the game by using '!deathcount set game <game>'"
        )
        return

    deaths = db_get_death_count(game, user.chatted_from)

    streamer = db_get_streamer_nickname(user.chatted_from)

    if deaths > 1:
        user.send_message(
            "{} has died {} times. RIP sfhSAD".format(streamer, deaths)
        )
    elif deaths == 1:
        user.send_message("{} has died 1 time. RIP sfhSAD".format(streamer))
    elif deaths == 0:
        user.send_message("{} has yet to die! sfhOH".format(streamer))
    else:
        user.send_message("{} is not in the list of games.".format(game))
        return

    logger.info(
        "Read death count of {} for the game {}, in the channel "
        "{}.".format(
            deaths, game, user.chatted_from, user.name
        )
    )

def add_game(user: IRC.User):
    game = add_game_regex.findall(user.message)[0]

    if not db_add_game(game, user.chatted_from):
        user.send_message("'{}' is already in the list of games.".format(game))
    else:
        logger.info(
            "{} has added the game '{}' to the channel '{}'".format(
                user.name, game, user.chatted_from
            )
        )
        user.send_message("'{}' is now in the list of games! sfhOH".format(game))

def set_game(user: IRC.User):
    game = set_game_regex.findall(user.message)[0]

    if db_set_default_game(game, user.chatted_from):
        logger.info(
            "Default game for channel '{}' has be set to '{}' by {}".format(
                user.chatted_from, game, user.name
            )
        )
        user.send_message("Game has been changed to {}. sfhOH".format(game))
    else:
        user.send_message(
            "{} is not in the list of games. If you'd like to add it to the "
            "list of games (and are a moderator), run '{}deathcount add game {}'."
            "".format(game, module_tools.DEFAULT_COMMAND_PREFIX, game)
        )

def remove_game(user: IRC.User):
    game = remove_game_regex.findall(user.message)[0]

    if db_remove_game(game, user.chatted_from):
        logger.info(
            "{} has removed the game '{}' from the channel '{}'".format(
                user.name, game, user.chatted_from
            )
        )
        user.send_message(
            "{} has been removed from the list of games! "
            "RIP".format(game)
        )

    else:
        user.send_message("{} is not in the list of games.".format(game))

def set_nickname(user: IRC.User):
    nickname = set_nickname_regex.findall(user.message)[0]

    db_set_nickname(nickname, user.chatted_from)
    logger.info(
        "{} has set the nickname for channel '{}' to '{}'".format(
            user.name, user.chatted_from, nickname
        )
    )
    user.send_message("Nickname has been set to '{}'".format(nickname))

# Registering Commands----------------------------------------------------------
module_tools.register_command("d( \d+)?", increment_deaths, "moderator",
                              mod_cooldown=15)
module_tools.register_command("dd( \d+)?", decrement_deaths, "moderator",
                              mod_cooldown=15)
module_tools.register_command("dc( \w+)?", death_count)
module_tools.register_command("deathcount reset", reset_deaths, "broadcaster")
module_tools.register_command("deathcount remove game (\w+|\W+)",
                              remove_game,
                              "moderator")
module_tools.register_command("deathcount add game (\w+|\W+)",
                              add_game,
                              "moderator",
                              mod_cooldown=15)
module_tools.register_command("deathcount set game (\w+|\W+)",
                              set_game,
                              "moderator")
module_tools.register_command("deathcount set nickname (\w+)",
                              set_nickname,
                              "moderator")
module_tools.register_shutdown_function(db_close_connection)

<commit_msg>Moved to seperate package/program<commit_after><|endoftext|>


<commit_before>import traceback
import re
import jedi

from .base import Base

class Source(Base):

    def __init__(self, vim):
        Base.__init__(self, vim)

        self.name = 'jedi'
        self.mark = '[jedi]'
        self.filetypes = ['python']
        self.rank = 1000
        self.min_pattern_length = 0
        self.input_pattern = r'[^. \t0-9]\.\w*|^\s*@\w*|^\s*from\s.+import \w*|^\s*from \w*|^\s*import \w*'
        self.is_bytepos = True

    def get_complete_position(self, context):
        return self.completions(1, 0)

    def gather_candidates(self, context):
        return self.completions(0, 0)

    def get_script(self, source=None, column=None):
        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.add_dot_after_module
        # Adds a dot after a module, because a module that is not accessed this way is definitely not the normal case.
        # However, in VIM this doesn’t work, that’s why it isn’t used at the moment.
        jedi.settings.add_dot_after_module = True

        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.add_bracket_after_function
        # Adds an opening bracket after a function, because that's normal behaviour.
        # Removed it again, because in VIM that is not very practical.
        jedi.settings.add_bracket_after_function = True

        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.additional_dynamic_modules
        # Additional modules in which Jedi checks if statements are to be found.
        # This is practical for IDEs, that want to administrate their modules themselves.
        jedi.settings.additional_dynamic_modules = \
            [b.name for b in self.vim.buffers if b.name is not None and b.name.endswith('.py')]

        # Need?
        if source is None:
            source = '\n'.join(self.vim.current.buffer)
        row = self.vim.current.window.cursor[0]
        # Need?
        if column is None:
            column = self.vim.current.window.cursor[1]
        buf_path = self.vim.current.buffer.name
        encoding = self.vim.eval('&encoding')

        return jedi.Script(source, row, column, buf_path, encoding)

    def completions(self, findstart, base):
        row, column = self.vim.current.window.cursor

        if findstart == 1:
            # Really good?
            return column
        else:
            # jedi-vim style? or simple?
            # source = '\n'.join(self.vim.current.buffer[:])
            source = ''
            for i, line in enumerate(self.vim.current.buffer):
                if i == row - 1:
                    source += line[:column] + str(base) + line[column:]
                else:
                    source += line
                source += '\n'

            script = self.get_script(source=source, column=column)
            completions = script.completions()

            out = []
            for c in completions:
                d = dict(word=c.complete,
                         abbr=c.name,
                         kind=c.description,
                         info=c.docstring(),
                         icase=1,
                         dup=1
                         )
                out.append(d)

            return out
<commit_msg>[WIP] Add omnifunc compatible? abbr description<commit_after>import traceback
import re
import jedi

from .base import Base

class Source(Base):

    def __init__(self, vim):
        Base.__init__(self, vim)

        self.name = 'jedi'
        self.mark = '[jedi]'
        self.filetypes = ['python']
        self.rank = 1000
        self.min_pattern_length = 0
        self.input_pattern = r'[^. \t0-9]\.\w*|^\s*@\w*|^\s*from\s.+import \w*|^\s*from \w*|^\s*import \w*'
        self.is_bytepos = True

    def get_complete_position(self, context):
        return self.completions(1, 0)

    def gather_candidates(self, context):
        return self.completions(0, 0)

    def get_script(self, source=None, column=None):
        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.add_dot_after_module
        # Adds a dot after a module, because a module that is not accessed this way is definitely not the normal case.
        # However, in VIM this doesn’t work, that’s why it isn’t used at the moment.
        jedi.settings.add_dot_after_module = True

        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.add_bracket_after_function
        # Adds an opening bracket after a function, because that's normal behaviour.
        # Removed it again, because in VIM that is not very practical.
        jedi.settings.add_bracket_after_function = True

        # http://jedi.jedidjah.ch/en/latest/docs/settings.html#jedi.settings.additional_dynamic_modules
        # Additional modules in which Jedi checks if statements are to be found.
        # This is practical for IDEs, that want to administrate their modules themselves.
        jedi.settings.additional_dynamic_modules = \
            [b.name for b in self.vim.buffers if b.name is not None and b.name.endswith('.py')]

        # Need?
        if source is None:
            source = '\n'.join(self.vim.current.buffer)
        row = self.vim.current.window.cursor[0]
        # Need?
        if column is None:
            column = self.vim.current.window.cursor[1]
        buf_path = self.vim.current.buffer.name
        encoding = self.vim.eval('&encoding')

        return jedi.Script(source, row, column, buf_path, encoding)

    def completions(self, findstart, base):
        row, column = self.vim.current.window.cursor

        if findstart == 1:
            # Really good?
            return column
        else:
            # jedi-vim style? or simple?
            # source = '\n'.join(self.vim.current.buffer[:])
            source = ''
            for i, line in enumerate(self.vim.current.buffer):
                if i == row - 1:
                    source += line[:column] + str(base) + line[column:]
                else:
                    source += line
                source += '\n'

            script = self.get_script(source=source, column=column)
            completions = script.completions()

            out = []
            for c in completions:
                if c.type == 'function':
                    abbr = re.sub('"(|)",', '', c.docstring().split("\n\n")[0].replace('\n', ' '))
                else:
                    abbr = c.name
                d = dict(word=c.complete,
                         abbr=str(abbr),
                         kind=c.description,
                         info=c.docstring(),
                         icase=1,
                         dup=1
                         )
                out.append(d)

            return out
<|endoftext|>