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()