use distinct for gridfs.list

Author: Mike Dirolf

gridfs/__init__.py

@@ -21,26 +21,36 @@
 """
 
 from gridfs.grid_file import GridFile
+from pymongo import (ASCENDING,
+                     DESCENDING)
 from pymongo.database import Database
 
 class GridFS(object):
     """An instance of GridFS on top of a single Database.
     """
-    def __init__(self, database):
+    def __init__(self, database, collection="fs"):
         """Create a new instance of :class:`GridFS`.
 
         Raises :class:`TypeError` if `database` is not an instance of
         :class:`~pymongo.database.Database`.
 
         :Parameters:
           - `database`: database to use
+          - `collection` (optional): root collection to use
+
+        .. versionadded:: 1.5.1+
+           The `collection` parameter.
 
         .. mongodoc:: gridfs
         """
         if not isinstance(database, Database):
             raise TypeError("database must be an instance of Database")
 
         self.__database = database
+        self.__files = database[collection].files
+        self.__chunks = database[collection].chunks
+        self.__chunks.create_index([("files_id", ASCENDING), ("n", ASCENDING)],
+                                   unique=True)
 
     def open(self, filename, mode="r", collection="fs"):
         """Open a :class:`~gridfs.grid_file.GridFile` for reading or
@@ -62,7 +72,8 @@ def open(self, filename, mode="r", collection="fs"):
           - `filename`: name of the :class:`~gridfs.grid_file.GridFile`
             to open
           - `mode` (optional): mode to open the file in
-          - `collection` (optional): root collection to use for this file
+          - `collection` (optional): root collection to use for this
+            file
         """
         return GridFile({"filename": filename}, self.__database, mode, collection)
 
@@ -116,7 +127,4 @@ def list(self, collection="fs"):
         """
         if not isinstance(collection, basestring):
             raise TypeError("collection must be an instance of basestring")
-        names = []
-        for grid_file in self.__database[collection].files.find():
-            names.append(grid_file["filename"])
-        return names
+        return self.__database[collection].files.distinct("filename")

test/test_gridfs.py

@@ -82,7 +82,8 @@ def test_list(self):
         f = self.fs.open("hello world", "w")
         f.close()
 
-        self.assertEqual(["mike", "test", "hello world"], self.fs.list())
+        self.assertEqual(set(["mike", "test", "hello world"]),
+                         set(self.fs.list()))
 
     def test_remove(self):
         self.assertRaises(TypeError, self.fs.remove, 5)
@@ -98,13 +99,14 @@ def test_remove(self):
         f = self.fs.open("hello world", "w")
         f.write("fly")
         f.close()
-        self.assertEqual(["mike", "test", "hello world"], self.fs.list())
+        self.assertEqual(set(["mike", "test", "hello world"]),
+                         set(self.fs.list()))
         self.assertEqual(self.db.fs.files.find().count(), 3)
         self.assertEqual(self.db.fs.chunks.find().count(), 3)
 
         self.fs.remove("test")
 
-        self.assertEqual(["mike", "hello world"], self.fs.list())
+        self.assertEqual(set(["mike", "hello world"]), set(self.fs.list()))
         self.assertEqual(self.db.fs.files.find().count(), 2)
         self.assertEqual(self.db.fs.chunks.find().count(), 2)
         f = self.fs.open("mike")
@@ -145,8 +147,8 @@ def test_list_alt_coll(self):
         f.close()
 
         self.assertEqual([], self.fs.list())
-        self.assertEqual(["mike", "test", "hello world"],
-                         self.fs.list("pymongo_test"))
+        self.assertEqual(set(["mike", "test", "hello world"]),
+                         set(self.fs.list("pymongo_test")))
 
     def test_remove_alt_coll(self):
         f = self.fs.open("mike", "w", "pymongo_test")
@@ -160,10 +162,11 @@ def test_remove_alt_coll(self):
         f.close()
 
         self.fs.remove("test")
-        self.assertEqual(["mike", "test", "hello world"],
-                         self.fs.list("pymongo_test"))
+        self.assertEqual(set(["mike", "test", "hello world"]),
+                         set(self.fs.list("pymongo_test")))
         self.fs.remove("test", "pymongo_test")
-        self.assertEqual(["mike", "hello world"], self.fs.list("pymongo_test"))
+        self.assertEqual(set(["mike", "hello world"]),
+                         set(self.fs.list("pymongo_test")))
 
         f = self.fs.open("mike", collection="pymongo_test")
         self.assertEqual(f.read(), "hi")
@@ -204,10 +207,10 @@ def test_threaded_writes(self):
         self.assertEqual(f.read(), "hello")
         f.close()
 
-    # NOTE I do recognize how gross this is. There is no good way to test the
-    # with statement because it is a syntax error in older python versions.
-    # One option would be to use eval and skip the test if it is a syntax
-    # error.
+    # NOTE I do recognize how gross this is. There is no good way to
+    # test the with statement because it is a syntax error in older
+    # python versions.  One option would be to use eval and skip the
+    # test if it is a syntax error.
     if sys.version_info[:2] == (2, 5):
         import gridfs15
         test_with_statement = gridfs15.test_with_statement