PYTHON-310: Added find() method to gridfs

3rf · 3rf · commit 237754dbefce · 2014-01-06T14:57:00.000-08:00
Allows iteration through arbitrary queries against the files collection using
a new GridOutCursor class.
diff --git a/doc/api/gridfs/grid_file.rst b/doc/api/gridfs/grid_file.rst
@@ -17,3 +17,6 @@
 
    .. autoclass:: GridFile
       :members:
+
+   .. autoclass:: GridOutCursor
+      :members:
diff --git a/gridfs/__init__.py b/gridfs/__init__.py
@@ -23,7 +23,8 @@
 from gridfs.errors import (NoFile,
                            UnsupportedAPI)
 from gridfs.grid_file import (GridIn,
-                              GridOut)
+                              GridOut,
+                              GridOutCursor)
 from pymongo import (MongoClient,
                      ASCENDING,
                      DESCENDING)
@@ -278,6 +279,72 @@ def list(self):
             name for name in self.__files.distinct("filename")
             if name is not None]
 
+    def find(self, *args, **kwargs):
+        """Query GridFS for files.
+
+        Returns a cursor that iterates across files matching
+        arbitrary queries on the files collection. Can be combined
+        with other modifiers for additional control. For example
+
+        >>> for grid_out in fs.find({"filename": "lisa.txt"}, timeout=False):
+        >>>     data = grid_out.read()
+
+        would iterate through all versions of "lisa.txt" stored in GridFS.
+        Note that setting timeout to False may be important to prevent the
+        cursor from timing out during long multi-file processing work.
+
+        As another example, the call
+
+        >>> most_recent_three = fs.find().sort("uploadDate", -1).limit(3)
+
+        would return a cursor to the three most recently uploaded files
+        in GridFS.
+
+        Follows a similar interface to
+        :meth:`~pymongo.collection.Collection.find`
+        in :class:`~pymongo.collection.Collection`.
+
+        :Parameters:
+          - `spec` (optional): a SON object specifying elements which
+            must be present for a document to be included in the
+            result set
+          - `skip` (optional): the number of files to omit (from
+            the start of the result set) when returning the results
+          - `limit` (optional): the maximum number of results to
+            return
+          - `timeout` (optional): if True (the default), any returned
+            cursor is closed by the server after 10 minutes of
+            inactivity. If set to False, the returned cursor will never
+            time out on the server. Care should be taken to ensure that
+            cursors with timeout turned off are properly closed.
+          - `sort` (optional): a list of (key, direction) pairs
+            specifying the sort order for this query. See
+            :meth:`~pymongo.cursor.Cursor.sort` for details.
+          - `max_scan` (optional): limit the number of file documents
+            examined when performing the query
+          - `read_preference` (optional): The read preference for
+            this query.
+          - `tag_sets` (optional): The tag sets for this query.
+          - `secondary_acceptable_latency_ms` (optional): Any replica-set
+            member whose ping time is within secondary_acceptable_latency_ms of
+            the nearest member may accept reads. Default 15 milliseconds.
+            **Ignored by mongos** and must be configured on the command line.
+            See the localThreshold_ option for more information.
+          - `compile_re` (optional): if ``False``, don't attempt to compile
+            BSON regex objects into Python regexes. Return instances of
+            :class:`~bson.regex.Regex` instead.
+
+        Raises :class:`TypeError` if any of the arguments are of
+        improper type. Returns an instance of
+        :class:`~gridfs.grid_file.GridOutCursor`
+        corresponding to this query.
+
+        .. versionadded:: 2.7
+        .. mongodoc:: find
+        .. _localThreshold: http://docs.mongodb.org/manual/reference/mongos/#cmdoption-mongos--localThreshold
+        """
+        return GridOutCursor(self.__collection, *args, **kwargs)
+
     def exists(self, document_or_id=None, **kwargs):
         """Check if a file exists in this instance of :class:`GridFS`.
 
diff --git a/gridfs/grid_file.py b/gridfs/grid_file.py
@@ -27,6 +27,7 @@
                            UnsupportedAPI)
 from pymongo import ASCENDING
 from pymongo.collection import Collection
+from pymongo.cursor import Cursor
 from pymongo.errors import DuplicateKeyError
 
 try:
@@ -614,3 +615,54 @@ class GridFile(object):
     def __init__(self, *args, **kwargs):
         raise UnsupportedAPI("The GridFile class is no longer supported. "
                              "Please use GridIn or GridOut instead.")
+
+
+class GridOutCursor(Cursor):
+    """A cursor / iterator for returning GridOut objects as the result
+    of an arbitrary query against the GridFS files collection.
+    """
+    def __init__(self, collection, spec=None, skip=0, limit=0,
+                 timeout=True, sort=None, max_scan=None,
+                 read_preference=None, tag_sets=None,
+                 secondary_acceptable_latency_ms=None, compile_re=True):
+        """Create a new cursor, similar to the normal
+        :class:`~pymongo.cursor.Cursor`.
+
+        Should not be called directly by application developers - see
+        the :class:`~gridfs.GridFS` method :meth:`~gridfs.GridFS.find` instead.
+
+        .. versionadded 2.7
+
+        .. mongodoc:: cursors
+        """
+        # Hold on to the base "fs" collection to create GridOut objects later.
+        self.__root_collection = collection
+
+        # Copy these settings from collection if they are not set by caller.
+        read_preference = read_preference or collection.files.read_preference
+        tag_sets = tag_sets or collection.files.tag_sets
+        latency = (secondary_acceptable_latency_ms
+                   or collection.files.secondary_acceptable_latency_ms)
+
+        super(GridOutCursor, self).__init__(
+            collection.files, spec, skip=skip, limit=limit, timeout=timeout,
+            sort=sort, max_scan=max_scan, read_preference=read_preference,
+            secondary_acceptable_latency_ms=latency, compile_re=compile_re,
+            tag_sets=tag_sets)
+
+    def next(self):
+        """Get next GridOut object from cursor.
+        """
+        next_file = super(GridOutCursor, self).next()
+        return GridOut(self.__root_collection, file_document=next_file)
+
+    def add_option(self, *args, **kwargs):
+        raise NotImplementedError("Method does not exist for GridOutCursor")
+
+    def remove_option(self, *args, **kwargs):
+        raise NotImplementedError("Method does not exist for GridOutCursor")
+
+    def _clone_base(self):
+        """Creates an empty GridOutCursor for information to be copied into.
+        """
+        return GridOutCursor(self.__root_collection)
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -219,11 +219,11 @@ def clone(self):
         unevaluated, even if the current instance has been partially or
         completely evaluated.
         """
-        return self.__clone(True)
+        return self._clone(True)
 
-    def __clone(self, deepcopy=True):
+    def _clone(self, deepcopy=True):
         self.__check_not_command_cursor('clone')
-        clone = Cursor(self.__collection)
+        clone = self._clone_base()
         values_to_clone = ("spec", "fields", "skip", "limit", "max_time_ms",
                            "comment", "max", "min",
                            "snapshot", "ordering", "explain", "hint",
@@ -235,10 +235,15 @@ def __clone(self, deepcopy=True):
         data = dict((k, v) for k, v in self.__dict__.iteritems()
                     if k.startswith('_Cursor__') and k[9:] in values_to_clone)
         if deepcopy:
-            data = self.__deepcopy(data)
+            data = self._deepcopy(data)
         clone.__dict__.update(data)
         return clone
 
+    def _clone_base(self):
+        """Creates an empty Cursor object for information to be copied into.
+        """
+        return Cursor(self.__collection)
+
     def __die(self):
         """Closes this cursor.
         """
@@ -1017,16 +1022,16 @@ def __copy__(self):
 
         .. versionadded:: 2.4
         """
-        return self.__clone(deepcopy=False)
+        return self._clone(deepcopy=False)
 
     def __deepcopy__(self, memo):
         """Support function for `copy.deepcopy()`.
 
         .. versionadded:: 2.4
         """
-        return self.__clone(deepcopy=True)
+        return self._clone(deepcopy=True)
 
-    def __deepcopy(self, x, memo=None):
+    def _deepcopy(self, x, memo=None):
         """Deepcopy helper for the data dictionary or list.
 
         Regular expressions cannot be deep copied but as they are immutable we
@@ -1046,7 +1051,7 @@ def __deepcopy(self, x, memo=None):
 
         for key, value in iterator:
             if isinstance(value, (dict, list)) and not isinstance(value, SON):
-                value = self.__deepcopy(value, memo)
+                value = self._deepcopy(value, memo)
             elif not isinstance(value, RE_TYPE):
                 value = copy.deepcopy(value, memo)
 
diff --git a/test/test_grid_file.py b/test/test_grid_file.py
@@ -32,7 +32,8 @@
                               _SEEK_END,
                               GridIn,
                               GridFile,
-                              GridOut)
+                              GridOut,
+                              GridOutCursor)
 from gridfs.errors import (NoFile,
                            UnsupportedAPI)
 from pymongo import MongoClient
@@ -589,6 +590,27 @@ def test_grid_in_lazy_connect(self):
         self.assertRaises(ConnectionFailure, infile.write, b('data goes here'))
         self.assertRaises(ConnectionFailure, infile.close)
 
+    def test_grid_out_cursor_options(self):
+        self.assertRaises(TypeError, GridOutCursor.__init__, self.db.fs, {},
+                                                    tailable=True)
+        self.assertRaises(TypeError, GridOutCursor.__init__, self.db.fs, {},
+                                                    fields={"filename":1})
+
+        cursor = GridOutCursor(self.db.fs, {})
+        min_ms = self.db.fs.files.secondary_acceptable_latency_ms
+        new_ms = cursor._Cursor__secondary_acceptable_latency_ms
+        self.assertEqual(min_ms, new_ms)
+        cursor = GridOutCursor(self.db.fs, {},
+                               secondary_acceptable_latency_ms=100)
+        min_ms = self.db.fs.files.secondary_acceptable_latency_ms
+        new_ms = cursor._Cursor__secondary_acceptable_latency_ms
+        self.assertNotEqual(min_ms, new_ms)
+        cursor_clone = cursor.clone()
+        self.assertEqual(cursor_clone.__dict__, cursor.__dict__)
+
+        self.assertRaises(NotImplementedError, cursor.add_option, 0)
+        self.assertRaises(NotImplementedError, cursor.remove_option, 0)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/test_gridfs.py b/test/test_gridfs.py
@@ -373,6 +373,22 @@ def test_gridfs_lazy_connect(self):
         f = fs.new_file()  # Still no connection.
         self.assertRaises(ConnectionFailure, f.close)
 
+    def test_gridfs_find(self):
+        self.fs.put(b("test2"), filename="two")
+        self.fs.put(b("test2+"), filename="two")
+        self.fs.put(b("test1"), filename="one")
+        self.fs.put(b("test2++"), filename="two")
+        self.assertEqual(3, self.fs.find({"filename":"two"}).count())
+        self.assertEqual(4, self.fs.find().count())
+        cursor = self.fs.find(timeout=False).sort("uploadDate", -1).skip(1).limit(2)
+        self.assertEqual(b("test1"), cursor.next().read())
+        cursor.rewind()
+        self.assertEqual(b("test1"), cursor.next().read())
+        self.assertEqual(b("test2+"), cursor.next().read())
+        self.assertRaises(StopIteration, cursor.next)
+        cursor.close()
+        self.assertRaises(TypeError, self.fs.find, {}, {"_id": True})
+
 
 class TestGridfsReplicaSet(TestReplicaSetClientBase):
     def test_gridfs_replica_set(self):
