Ensure get_or_create() works with new GFK API, refs coleifer#313

peewee.py

@@ -4501,9 +4501,15 @@ def get(cls, *query, **kwargs):
     @classmethod
     def get_or_create(cls, **kwargs):
         defaults = kwargs.pop('defaults', {})
-        sq = cls.select().filter(**kwargs)
+        query = cls.select()
+        for field, value in kwargs.items():
+            if '__' in field:
+                query = query.filter(**{field: value})
+            else:
+                query = query.where(getattr(cls, field) == value)
+
         try:
-            return sq.get(), False
+            return query.get(), False
         except cls.DoesNotExist:
             try:
                 params = dict((k, v) for k, v in kwargs.items()
@@ -4513,7 +4519,7 @@ def get_or_create(cls, **kwargs):
                     return cls.create(**params), True
             except IntegrityError as exc:
                 try:
-                    return sq.get(), False
+                    return query.get(), False
                 except cls.DoesNotExist:
                     raise exc
 

playhouse/tests/test_gfk.py

@@ -86,6 +86,33 @@ def test_querying(self):
         tags = Tag.select().where(Tag.object == tacos).order_by(Tag.tag)
         self.assertEqual([tag.tag for tag in tags], ['fried', 'spicy'])
 
+    def test_get_or_create(self):
+        a = Appetizer.create(name='walrus mix')
+        tag, created = Tag.get_or_create(tag='walrus-food', object=a)
+        self.assertTrue(created)
+        self.assertEqual(tag.object, a)
+
+        tag_db = Tag.get(Tag.id == tag.id)
+        self.assertEqual(tag_db.object, a)
+
+        tag, created = Tag.get_or_create(tag='walrus-food', object=a)
+        self.assertFalse(created)
+        self.assertEqual(Tag.select().count(), 1)
+        self.assertEqual(tag, tag_db)
+
+        tag2, created = Tag.get_or_create(tag='walrus-treats', object=a)
+        self.assertTrue(created)
+        tag2_db = Tag.get(Tag.id == tag2.id)
+        self.assertEqual(tag2_db.tag, 'walrus-treats')
+        self.assertEqual(tag2_db.object, a)
+
+        b = Appetizer.create(name='walrus-meal')
+        tag3, created = Tag.get_or_create(tag='walrus-treats', object=b)
+        self.assertTrue(created)
+        tag3_db = Tag.get(Tag.id == tag3.id)
+        self.assertEqual(tag3_db.tag, 'walrus-treats')
+        self.assertEqual(tag3_db.object, b)
+
     def test_gfk_api(self):
         self.create()