Added SequenceField.set_next_value(value) helper (#159)

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 
 Changes in 0.8.X
 ================
+- Added SequenceField.set_next_value(value) helper (#159)
 - Updated .only() behaviour - now like exclude it is chainable (#202)
 - Added with_limit_and_skip support to count() (#235)
 - Removed __len__ from queryset (#247)

mongoengine/fields.py

@@ -1408,13 +1408,13 @@ class SequenceField(BaseField):
     COLLECTION_NAME = 'mongoengine.counters'
     VALUE_DECORATOR = int
 
-    def __init__(self, collection_name=None, db_alias=None,
+    def __init__(self, collection_name=None, db_alias=None, sequence_name=None,
-            sequence_name=None, value_decorator=None, *args, **kwargs):
+                 value_decorator=None, *args, **kwargs):
         self.collection_name = collection_name or self.COLLECTION_NAME
         self.db_alias = db_alias or DEFAULT_CONNECTION_NAME
         self.sequence_name = sequence_name
         self.value_decorator = (callable(value_decorator) and
-            value_decorator or self.VALUE_DECORATOR)
+                                value_decorator or self.VALUE_DECORATOR)
         return super(SequenceField, self).__init__(*args, **kwargs)
 
     def generate(self):
@@ -1430,6 +1430,17 @@ class SequenceField(BaseField):
                                              upsert=True)
         return self.value_decorator(counter['next'])
 
+    def set_next_value(self, value):
+        """Helper method to set the next sequence value"""
+        sequence_name = self.get_sequence_name()
+        sequence_id = "%s.%s" % (sequence_name, self.name)
+        collection = get_db(alias=self.db_alias)[self.collection_name]
+        counter = collection.find_and_modify(query={"_id": sequence_id},
+                                             update={"$set": {"next": value}},
+                                             new=True,
+                                             upsert=True)
+        return self.value_decorator(counter['next'])
+
     def get_sequence_name(self):
         if self.sequence_name:
             return self.sequence_name
@@ -1438,7 +1449,7 @@ class SequenceField(BaseField):
             return owner._get_collection_name()
         else:
             return ''.join('_%s' % c if c.isupper() else c
-                            for c in owner._class_name).strip('_').lower()
+                           for c in owner._class_name).strip('_').lower()
 
     def __get__(self, instance, owner):
         value = super(SequenceField, self).__get__(instance, owner)

tests/fields/fields.py

@@ -2164,18 +2164,21 @@ class FieldTest(unittest.TestCase):
         Person.drop_collection()
 
         for x in xrange(10):
-            p = Person(name="Person %s" % x)
+            Person(name="Person %s" % x).save()
-            p.save()
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
         self.assertEqual(c['next'], 10)
 
         ids = [i.id for i in Person.objects]
-        self.assertEqual(ids, range(1, 11))
+        self.assertEqual(ids, xrange(1, 11))
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
         self.assertEqual(c['next'], 10)
 
+        Person.id.set_next_value(1000)
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 1000)
+
     def test_sequence_field_sequence_name(self):
         class Person(Document):
             id = SequenceField(primary_key=True, sequence_name='jelly')
@@ -2185,8 +2188,7 @@ class FieldTest(unittest.TestCase):
         Person.drop_collection()
 
         for x in xrange(10):
-            p = Person(name="Person %s" % x)
+            Person(name="Person %s" % x).save()
-            p.save()
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'})
         self.assertEqual(c['next'], 10)
@@ -2197,6 +2199,10 @@ class FieldTest(unittest.TestCase):
         c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'})
         self.assertEqual(c['next'], 10)
 
+        Person.id.set_next_value(1000)
+        c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'})
+        self.assertEqual(c['next'], 1000)
+
     def test_multiple_sequence_fields(self):
         class Person(Document):
             id = SequenceField(primary_key=True)
@@ -2207,21 +2213,28 @@ class FieldTest(unittest.TestCase):
         Person.drop_collection()
 
         for x in xrange(10):
-            p = Person(name="Person %s" % x)
+            Person(name="Person %s" % x).save()
-            p.save()
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
         self.assertEqual(c['next'], 10)
 
         ids = [i.id for i in Person.objects]
-        self.assertEqual(ids, range(1, 11))
+        self.assertEqual(ids, xrange(1, 11))
 
         counters = [i.counter for i in Person.objects]
-        self.assertEqual(counters, range(1, 11))
+        self.assertEqual(counters, xrange(1, 11))
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
         self.assertEqual(c['next'], 10)
 
+        Person.id.set_next_value(1000)
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 1000)
+
+        Person.counter.set_next_value(999)
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.counter'})
+        self.assertEqual(c['next'], 999)
+
     def test_sequence_fields_reload(self):
         class Animal(Document):
             counter = SequenceField()
@@ -2230,8 +2243,7 @@ class FieldTest(unittest.TestCase):
         self.db['mongoengine.counters'].drop()
         Animal.drop_collection()
 
-        a = Animal(name="Boi")
+        a = Animal(name="Boi").save()
-        a.save()
 
         self.assertEqual(a.counter, 1)
         a.reload()
@@ -2261,10 +2273,8 @@ class FieldTest(unittest.TestCase):
         Person.drop_collection()
 
         for x in xrange(10):
-            a = Animal(name="Animal %s" % x)
+            Animal(name="Animal %s" % x).save()
-            a.save()
+            Person(name="Person %s" % x).save()
-            p = Person(name="Person %s" % x)
-            p.save()
 
         c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
         self.assertEqual(c['next'], 10)