<commit_before>from peewee import *

from .base import ModelTestCase
from .base import TestModel
from .base import get_in_memory_db
from .base import requires_models
from .base_models import Tweet
from .base_models import User


class User(TestModel):
    username = TextField(unique=True)


class Note(TestModel):
    text = TextField()
    users = ManyToManyField(User)


NoteUserThrough = Note.users.get_through_model()

AltThroughDeferred = DeferredThroughModel()


class AltNote(TestModel):
    text = TextField()
    users = ManyToManyField(User, through_model=AltThroughDeferred)


class AltThroughModel(TestModel):
    user = ForeignKeyField(User, backref='_xx_rel')
    note = ForeignKeyField(AltNote, backref='_xx_rel')

    class Meta:
        primary_key = CompositeKey('user', 'note')


AltThroughDeferred.set_model(AltThroughModel)


class TestManyToMany(ModelTestCase):
    database = get_in_memory_db()
    requires = [User, Note, NoteUserThrough, AltNote, AltThroughModel]

    user_to_note = {
        'gargie': [1, 2],
        'huey': [2, 3],
        'mickey': [3, 4],
        'zaizee': [4, 5],
    }

    def setUp(self):
        super(TestManyToMany, self).setUp()
        for username in sorted(self.user_to_note):
            User.create(username=username)
        for i in range(5):
            Note.create(text='note-%s' % (i + 1))

    def test_through_model(self):
        self.assertEqual(len(NoteUserThrough._meta.fields), 3)
        fields = NoteUserThrough._meta.fields
        self.assertEqual(sorted(fields), ['id', 'note', 'user'])

        note_field = fields['note']
        self.assertEqual(note_field.rel_model, Note)
        self.assertFalse(note_field.null)

        user_field = fields['user']
        self.assertEqual(user_field.rel_model, User)
        self.assertFalse(user_field.null)

    def _set_data(self):
        for username, notes in self.user_to_note.items():
            user = User.get(User.username == username)
            for note in notes:
                NoteUserThrough.create(
                    note=Note.get(Note.text == 'note-%s' % note),
                    user=user)

    def assertNotes(self, query, expected):
        notes = [note.text for note in query]
        self.assertEqual(sorted(notes),
                         ['note-%s' % i for i in sorted(expected)])

    def assertUsers(self, query, expected):
        usernames = [user.username for user in query]
        self.assertEqual(sorted(usernames), sorted(expected))

    def test_accessor_query(self):
        self._set_data()
        gargie, huey, mickey, zaizee = User.select().order_by(User.username)

        with self.assertQueryCount(1):
            self.assertNotes(gargie.notes, [1, 2])
        with self.assertQueryCount(1):
            self.assertNotes(zaizee.notes, [4, 5])
        with self.assertQueryCount(2):
            self.assertNotes(User.create(username='x').notes, [])

        n1, n2, n3, n4, n5 = Note.select().order_by(Note.text)
        with self.assertQueryCount(1):
            self.assertUsers(n1.users, ['gargie'])
        with self.assertQueryCount(1):
            self.assertUsers(n2.users, ['gargie', 'huey'])
        with self.assertQueryCount(1):
            self.assertUsers(n5.users, ['zaizee'])
        with self.assertQueryCount(2):
            self.assertUsers(Note.create(text='x').users, [])

    def test_prefetch_notes(self):
        self._set_data()
        gargie, huey, mickey, zaizee = prefetch(User.select().order_by(User.username), NoteUserThrough, Note)

        with self.assertQueryCount(0):
            self.assertNotes(gargie.notes, [1, 2])
        with self.assertQueryCount(0):
            self.assertNotes(zaizee.notes, [4, 5])
        with self.assertQueryCount(2):
            self.assertNotes(User.create(username='x').notes, [])

    def test_prefetch_users(self):
        self._set_data()
        n1, n2, n3, n4, n5 = prefetch(Note.select().order_by(Note.text), NoteUserThrough, User)

        with self.assertQueryCount(0):
            self.assertUsers(n1.users, ['gargie'])
        with self.assertQueryCount(0):
            self.assertUsers(n2.users, ['gargie', 'huey'])
        with self.assertQueryCount(0):
            self.assertUsers(n5.users, ['zaizee'])
        with self.assertQueryCount(2):
            self.assertUsers(Note.create(text='x').users, [])

    def test_query_filtering(self):
        self._set_data()
        gargie, huey, mickey, zaizee = User.select().order_by(User.username)

        with self.assertQueryCount(1):
            notes = gargie.notes.where(Note.text != 'note-2')
            self.assertNotes(notes, [1])

    def test_set_value(self):
        self._set_data()<commit_msg>Additional assertion, pep8 for manytomany prefetch tests.<commit_after>from peewee import *

from .base import ModelTestCase
from .base import TestModel
from .base import get_in_memory_db
from .base import requires_models
from .base_models import Tweet
from .base_models import User


class User(TestModel):
    username = TextField(unique=True)


class Note(TestModel):
    text = TextField()
    users = ManyToManyField(User)


NoteUserThrough = Note.users.get_through_model()

AltThroughDeferred = DeferredThroughModel()


class AltNote(TestModel):
    text = TextField()
    users = ManyToManyField(User, through_model=AltThroughDeferred)


class AltThroughModel(TestModel):
    user = ForeignKeyField(User, backref='_xx_rel')
    note = ForeignKeyField(AltNote, backref='_xx_rel')

    class Meta:
        primary_key = CompositeKey('user', 'note')


AltThroughDeferred.set_model(AltThroughModel)


class TestManyToMany(ModelTestCase):
    database = get_in_memory_db()
    requires = [User, Note, NoteUserThrough, AltNote, AltThroughModel]

    user_to_note = {
        'gargie': [1, 2],
        'huey': [2, 3],
        'mickey': [3, 4],
        'zaizee': [4, 5],
    }

    def setUp(self):
        super(TestManyToMany, self).setUp()
        for username in sorted(self.user_to_note):
            User.create(username=username)
        for i in range(5):
            Note.create(text='note-%s' % (i + 1))

    def test_through_model(self):
        self.assertEqual(len(NoteUserThrough._meta.fields), 3)
        fields = NoteUserThrough._meta.fields
        self.assertEqual(sorted(fields), ['id', 'note', 'user'])

        note_field = fields['note']
        self.assertEqual(note_field.rel_model, Note)
        self.assertFalse(note_field.null)

        user_field = fields['user']
        self.assertEqual(user_field.rel_model, User)
        self.assertFalse(user_field.null)

    def _set_data(self):
        for username, notes in self.user_to_note.items():
            user = User.get(User.username == username)
            for note in notes:
                NoteUserThrough.create(
                    note=Note.get(Note.text == 'note-%s' % note),
                    user=user)

    def assertNotes(self, query, expected):
        notes = [note.text for note in query]
        self.assertEqual(sorted(notes),
                         ['note-%s' % i for i in sorted(expected)])

    def assertUsers(self, query, expected):
        usernames = [user.username for user in query]
        self.assertEqual(sorted(usernames), sorted(expected))

    def test_accessor_query(self):
        self._set_data()
        gargie, huey, mickey, zaizee = User.select().order_by(User.username)

        with self.assertQueryCount(1):
            self.assertNotes(gargie.notes, [1, 2])
        with self.assertQueryCount(1):
            self.assertNotes(zaizee.notes, [4, 5])
        with self.assertQueryCount(2):
            self.assertNotes(User.create(username='x').notes, [])

        n1, n2, n3, n4, n5 = Note.select().order_by(Note.text)
        with self.assertQueryCount(1):
            self.assertUsers(n1.users, ['gargie'])
        with self.assertQueryCount(1):
            self.assertUsers(n2.users, ['gargie', 'huey'])
        with self.assertQueryCount(1):
            self.assertUsers(n5.users, ['zaizee'])
        with self.assertQueryCount(2):
            self.assertUsers(Note.create(text='x').users, [])

    def test_prefetch_notes(self):
        self._set_data()
        with self.assertQueryCount(3):
            gargie, huey, mickey, zaizee = prefetch(
                User.select().order_by(User.username),
                NoteUserThrough,
                Note)

        with self.assertQueryCount(0):
            self.assertNotes(gargie.notes, [1, 2])
        with self.assertQueryCount(0):
            self.assertNotes(zaizee.notes, [4, 5])
        with self.assertQueryCount(2):
            self.assertNotes(User.create(username='x').notes, [])

    def test_prefetch_users(self):
        self._set_data()
        with self.assertQueryCount(3):
            n1, n2, n3, n4, n5 = prefetch(
                Note.select().order_by(Note.text),
                NoteUserThrough,
                User)

        with self.assertQueryCount(0):
            self.assertUsers(n1.users, ['gargie'])
        with self.assertQueryCount(0):
            self.assertUsers(n2.users, ['gargie', 'huey'])
        with self.assertQueryCount(0):
            self.assertUsers(n5.users, ['zaizee'])
        with self.assertQueryCount(2):
            self.assertUsers(Note.create(text='x').users, [])

    def test_query_filtering(self):
        self._set_data()
        gargie, huey, mickey, zaizee = User.select().order_by(User.username)

        with self.assertQueryCount(1):
            notes = gargie.notes.where(Note.text != 'note-2')
            self.assertNotes(notes, [1])

    def test_set_value(self):
        self._set_data()
{'fraction_non_alphanumeric': 0.12528521053723293, 'fraction_numerical': 0.01109728272142709, 'mean_word_length': 2.655420773313116, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '616141', 'n_tokens_mistral': 3119, 'n_tokens_neox': 2982, 'n_words': 571}