consolidate index document generation

pymongo/__init__.py

@@ -14,6 +14,10 @@
 
 """A Mongo driver for Python."""
 
+import types
+
+from pymongo.son import SON
+
 ASCENDING = 1
 """Ascending sort order."""
 DESCENDING = -1
@@ -26,3 +30,34 @@ SLOW_ONLY = 1
 ALL = 2
 """Profile all operations."""
 
+def _index_list(key_or_list, direction):
+    """Helper to generate a list of (key, direction) pairs.
+
+    Takes such a list, or a single key and direction.
+    """
+    if direction is not None:
+        return [(key_or_list, direction)]
+    else:
+        return key_or_list
+
+def _index_document(index_list):
+    """Helper to generate an index specifying document.
+
+    Takes a list of (key, direction) pairs.
+    """
+    if not isinstance(index_list, types.ListType):
+        raise TypeError("if no direction is specified, key_or_list must be an"
+                        "instance of list")
+    if not len(index_list):
+        raise ValueError("key_or_list must not be the empty list")
+
+    index = SON()
+    for (key, value) in index_list:
+        if not isinstance(key, types.StringTypes):
+            raise TypeError("first item in each key pair must be a string")
+        if not isinstance(value, types.IntType):
+            raise TypeError("second item in each key pair must be ASCENDING or"
+                            "DESCENDING")
+        index[key] = value
+    return index
+

pymongo/collection.py

@@ -16,6 +16,7 @@
 
 import types
 
+import pymongo
 import bson
 from objectid import ObjectId
 from cursor import Cursor
@@ -312,28 +313,11 @@ class Collection(object):
           - `direction` (optional): must be included if key_or_list is a single
             key, otherwise must be None
         """
-        if direction:
-            keys = [(key_or_list, direction)]
-        else:
-            keys = key_or_list
-
-        if not isinstance(keys, types.ListType):
-            raise TypeError("if no direction is specified, key_or_list must be an instance of list")
-        if not len(keys):
-            raise ValueError("key_or_list must not be the empty list")
-
         to_save = SON()
+        keys = pymongo._index_list(key_or_list, direction)
         to_save["name"] = self._gen_index_name(keys)
         to_save["ns"] = self.full_name()
-
-        key_object = SON()
-        for (key, value) in keys:
-            if not isinstance(key, types.StringTypes):
-                raise TypeError("first item in each key pair must be a string")
-            if not isinstance(value, types.IntType):
-                raise TypeError("second item in each key pair must be ASCENDING or DESCENDING")
-            key_object[key] = value
-        to_save["key"] = key_object
+        to_save["key"] = pymongo._index_document(keys)
 
         self.__database.system.indexes.save(to_save, False)
         return to_save["name"]

pymongo/cursor.py

@@ -18,6 +18,7 @@ import types
 import struct
 from threading import Lock
 
+import pymongo
 import bson
 from son import SON
 from code import Code
@@ -140,27 +141,8 @@ class Cursor(object):
             key, otherwise must be None
         """
         self.__check_okay_to_chain()
-
-        # TODO a lot of this logic could be shared with create_index()
-        if direction:
-            keys = [(key_or_list, direction)]
-        else:
-            keys = key_or_list
-
-        if not isinstance(keys, types.ListType):
-            raise TypeError("if no direction is specified, key_or_list must be an instance of list")
-        if not len(keys):
-            raise ValueError("key_or_list must not be the empty list")
-
-        orderby = SON()
-        for (key, value) in keys:
-            if not isinstance(key, types.StringTypes):
-                raise TypeError("first item in each key pair must be a string")
-            if not isinstance(value, types.IntType):
-                raise TypeError("second item in each key pair must be ASCENDING or DESCENDING")
-            orderby[key] = value
-
-        self.__ordering = orderby
+        keys = pymongo._index_list(key_or_list, direction)
+        self.__ordering = pymongo._index_document(keys)
         return self
 
     def count(self):