Cursor add_option/remove_option PYTHON-242

You can now set and unset arbitrary query flags.

Set the tailable flag:
cursor = db.coll.find().add_option(2)

Now unset it:
cursor.remove_option(2)

Author: behackett

pymongo/cursor.py

@@ -112,6 +112,7 @@ def __init__(self, collection, spec=None, fields=None, skip=0, limit=0,
         self.__tz_aware = collection.database.connection.tz_aware
         self.__must_use_master = _must_use_master
         self.__is_command = _is_command
+        self.__query_flags = 0
 
         self.__data = []
         self.__connection_id = None
@@ -174,6 +175,7 @@ def clone(self):
         copy.__partial = self.__partial
         copy.__must_use_master = self.__must_use_master
         copy.__is_command = self.__is_command
+        copy.__query_flags = self.__query_flags
         copy.__kwargs = self.__kwargs
         return copy
 
@@ -209,7 +211,7 @@ def __query_spec(self):
     def __query_options(self):
         """Get the query options string to use for this query.
         """
-        options = 0
+        options = self.__query_flags
         if self.__tailable:
             options |= _QUERY_OPTIONS["tailable_cursor"]
         if self.__slave_okay:
@@ -228,6 +230,32 @@ def __check_okay_to_chain(self):
         if self.__retrieved or self.__id is not None:
             raise InvalidOperation("cannot set options after executing query")
 
+    def add_option(self, mask):
+        """Set arbitary query flags using a bitmask.
+        
+        To set the tailable flag:
+        cursor.add_option(2)
+        """
+        if not isinstance(mask, int):
+            raise TypeError("mask must be an int")
+        self.__check_okay_to_chain()
+
+        self.__query_flags |= mask
+        return self
+
+    def remove_option(self, mask):
+        """Unset arbitrary query flags using a bitmask.
+        
+        To unset the tailable flag:
+        cursor.remove_option(2)
+        """
+        if not isinstance(mask, int):
+            raise TypeError("mask must be an int")
+        self.__check_okay_to_chain()
+
+        self.__query_flags &= ~mask
+        return self
+
     def limit(self, limit):
         """Limits the number of results to be returned by this cursor.
 

test/test_cursor.py

@@ -464,16 +464,62 @@ class MyClass(dict):
                                    slave_okay=True,
                                    await_data=True,
                                    partial=True).limit(2)
+        cursor.add_option(64)
 
         cursor2 = cursor.clone()
         self.assertEqual(cursor._Cursor__skip, cursor2._Cursor__skip)
+        self.assertEqual(cursor._Cursor__limit, cursor2._Cursor__limit)
         self.assertEqual(cursor._Cursor__timeout, cursor2._Cursor__timeout)
         self.assertEqual(cursor._Cursor__snapshot, cursor2._Cursor__snapshot)
         self.assertEqual(cursor._Cursor__tailable, cursor2._Cursor__tailable)
-        self.assertEqual(type(cursor._Cursor__as_class), type(cursor2._Cursor__as_class))
-        self.assertEqual(cursor._Cursor__slave_okay, cursor2._Cursor__slave_okay)
-        self.assertEqual(cursor._Cursor__await_data, cursor2._Cursor__await_data)
+        self.assertEqual(type(cursor._Cursor__as_class),
+                         type(cursor2._Cursor__as_class))
+        self.assertEqual(cursor._Cursor__slave_okay,
+                         cursor2._Cursor__slave_okay)
+        self.assertEqual(cursor._Cursor__await_data,
+                         cursor2._Cursor__await_data)
         self.assertEqual(cursor._Cursor__partial, cursor2._Cursor__partial)
+        self.assertEqual(cursor._Cursor__query_flags,
+                         cursor2._Cursor__query_flags)
+
+    def test_add_remove_option(self):
+        cursor = self.db.test.find()
+        self.assertEqual(0, cursor._Cursor__query_options())
+        cursor.add_option(2)
+        cursor2 = self.db.test.find(tailable=True)
+        self.assertEqual(2, cursor2._Cursor__query_options())
+        self.assertEqual(cursor._Cursor__query_options(),
+                         cursor2._Cursor__query_options())
+        cursor.add_option(32)
+        cursor2 = self.db.test.find(tailable=True, await_data=True)
+        self.assertEqual(34, cursor2._Cursor__query_options())
+        self.assertEqual(cursor._Cursor__query_options(),
+                         cursor2._Cursor__query_options())
+        cursor.add_option(128)
+        cursor2 = self.db.test.find(tailable=True,
+                                    await_data=True).add_option(128)
+        self.assertEqual(162, cursor2._Cursor__query_options())
+        self.assertEqual(cursor._Cursor__query_options(),
+                         cursor2._Cursor__query_options())
+
+        self.assertEqual(162, cursor._Cursor__query_options())
+        cursor.add_option(128)
+        self.assertEqual(162, cursor._Cursor__query_options())
+
+        cursor.remove_option(128)
+        cursor2 = self.db.test.find(tailable=True, await_data=True)
+        self.assertEqual(34, cursor2._Cursor__query_options())
+        self.assertEqual(cursor._Cursor__query_options(),
+                         cursor2._Cursor__query_options())
+        cursor.remove_option(32)
+        cursor2 = self.db.test.find(tailable=True)
+        self.assertEqual(2, cursor2._Cursor__query_options())
+        self.assertEqual(cursor._Cursor__query_options(),
+                         cursor2._Cursor__query_options())
+
+        self.assertEqual(2, cursor._Cursor__query_options())
+        cursor.remove_option(32)
+        self.assertEqual(2, cursor._Cursor__query_options())
 
     def test_count_with_fields(self):
         self.db.test.drop()