Cursors now support setting comments

hensom · A. Jesse Jiryu Davis · commit c936e182afcc · 2013-12-14T13:01:45.000-05:00
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -145,6 +145,7 @@ def __init__(self, collection, spec=None, fields=None, skip=0, limit=0,
         self.__max_scan = max_scan
         self.__explain = False
         self.__hint = None
+        self.__comment = None
         self.__as_class = as_class
         self.__slave_okay = slave_okay
         self.__manipulate = manipulate
@@ -271,6 +272,8 @@ def __query_spec(self):
             operators["$explain"] = True
         if self.__hint:
             operators["$hint"] = self.__hint
+        if self.__comment:
+            operators["$comment"] = self.__comment
         if self.__snapshot:
             operators["$snapshot"] = True
         if self.__max_scan:
@@ -657,6 +660,8 @@ def count(self, with_limit_and_skip=False):
         command['_use_master'] = use_master
         if self.__max_time_ms is not None:
             command["maxTimeMS"] = self.__max_time_ms
+        if self.__comment:
+            command['$comment'] = self.__comment
 
         if with_limit_and_skip:
             if self.__limit:
@@ -715,6 +720,8 @@ def distinct(self, key):
         options['_use_master'] = use_master
         if self.__max_time_ms is not None:
             options['maxTimeMS'] = self.__max_time_ms
+        if self.__comment:
+            options['$comment'] = self.__comment
 
         database = self.__collection.database
         return database.command("distinct",
@@ -766,6 +773,17 @@ def hint(self, index):
         self.__hint = helpers._index_document(index)
         return self
 
+    def comment(self, comment):
+        """Adds a 'comment', to the cursor.
+
+        http://docs.mongodb.org/manual/reference/operator/comment/
+
+        :Parameters:
+          - `comment`: A string or document
+        """
+        self.__comment = comment
+        return self
+
     def where(self, code):
         """Adds a $where clause to this query.
 
diff --git a/test/test_cursor.py b/test/test_cursor.py
@@ -28,7 +28,9 @@
 from bson.code import Code
 from bson.son import SON
 from pymongo import (ASCENDING,
-                     DESCENDING)
+                     DESCENDING,
+                     ALL,
+                     OFF)
 from pymongo.database import Database
 from pymongo.errors import (InvalidOperation,
                             OperationFailure,
@@ -38,6 +40,7 @@
 from test import version
 from test.test_client import get_client
 from test.utils import is_mongos, get_command_line
+from contextlib import contextmanager
 
 
 class TestCursor(unittest.TestCase):
@@ -939,5 +942,37 @@ def test_with_statement(self):
 """
         self.assertTrue(c1.alive)
 
+    def test_comment(self):
+        @contextmanager
+        def new_profile():
+            self.db.set_profiling_level(OFF)
+            self.db.system.profile.drop()
+            self.db.set_profiling_level(ALL)
+            yield
+            self.db.set_profiling_level(OFF)
+
+        with new_profile():
+            matching = list(self.db.test.find({'type': 'string'}).comment('foo'))
+            op = self.db.system.profile.find({'ns': 'pymongo_test.test',
+                                              'op': 'query',
+                                              'query.$comment': 'foo'})
+            self.assertEqual(op.count(), 1)
+
+        with new_profile():
+            self.db.test.find({'type': 'string'}).comment('foo').count()
+            op = self.db.system.profile.find({'ns': 'pymongo_test.$cmd',
+                                              'op': 'command',
+                                              'command.count': 'test',
+                                              'command.$comment': 'foo'})
+            self.assertEqual(op.count(), 1)
+
+        with new_profile():
+            self.db.test.find({'type': 'string'}).comment('foo').distinct('type')
+            op = self.db.system.profile.find({'ns': 'pymongo_test.$cmd',
+                                              'op': 'command',
+                                              'command.distinct': 'test',
+                                              'command.$comment': 'foo'})
+            self.assertEqual(op.count(), 1)
+
 if __name__ == "__main__":
     unittest.main()
