Fix Peewee Model.get() (cf. PR #1202)

1 files changed, 25 insertions, 1 deletions

diff --git a/archivist/peewee_ext.py b/archivist/peewee_ext.py
index 80bb6f4..a5a360b 100644
--- a/archivist/peewee_ext.py
+++ b/archivist/peewee_ext.py
@@ -1,5 +1,5 @@
 from playhouse.sqlite_ext import VirtualModel, VirtualIntegerField, VirtualCharField
-from peewee import Field
+from peewee import Field, OP, DJANGO_MAP, ForeignKeyField, ReverseRelationDescriptor, Expression, Query
 
 from itertools import starmap
 from functools import reduce
@@ -11,6 +11,30 @@ def sqlite_tuple_in(fields, values):
     subqueries = (reduce(op.and_, starmap(op.eq, zip(fields, value_tuple))) for value_tuple in values)
     return reduce(op.or_, subqueries)
 
+def convert_dict_to_node(self, qdict):
+    accum = []
+    joins = []
+    relationship = (ForeignKeyField, ReverseRelationDescriptor)
+    for key, value in sorted(qdict.items()):
+        curr = self.model_class
+        if '__' in key and key.rsplit('__', 1)[1] in DJANGO_MAP:
+            key, op = key.rsplit('__', 1)
+            op = DJANGO_MAP[op]
+        elif value is None:
+            op = OP.IS
+        else:
+            op = OP.EQ
+
+        for piece in key.split('__'):
+            model_attr = getattr(curr, piece)
+            if value is not None and isinstance(model_attr, relationship):
+                curr = model_attr.rel_model
+                joins.append(model_attr)
+        accum.append(Expression(model_attr, op, value))
+    return accum, joins
+
+Query.convert_dict_to_node = convert_dict_to_node
+
 class EnumField(Field):
     db_field = 'enum'