PYTHON-3524 Support passing list of strings to create_index (#1121)

Author: blink1073

doc/changelog.rst

@@ -1,6 +1,20 @@
 Changelog
 =========
 
+Changes in Version 4.4
+-----------------------
+
+- Added support for passing a list containing (key, direction) pairs
+  or keys to :meth:`~pymongo.collection.Collection.create_index`.
+
+Issues Resolved
+...............
+
+See the `PyMongo 4.4 release notes in JIRA`_ for the list of resolved issues
+in this release.
+
+.. _PyMongo 4.4 release notes in JIRA: https://jira.mongodb.org/secure/ReleaseNote.jspa?projectId=10004&version=34354
+
 Changes in Version 4.3.3
 ------------------------
 

pymongo/collection.py

@@ -26,7 +26,6 @@
     NoReturn,
     Optional,
     Sequence,
-    Tuple,
     Union,
 )
 
@@ -62,6 +61,8 @@
     ReplaceOne,
     UpdateMany,
     UpdateOne,
+    _IndexKeyHint,
+    _IndexList,
 )
 from pymongo.read_preferences import ReadPreference, _ServerMode
 from pymongo.results import (
@@ -85,9 +86,6 @@
     UpdateOne,
     UpdateMany,
 ]
-# Hint supports index name, "myIndex", or list of index pairs: [('x', 1), ('y', -1)]
-_IndexList = Sequence[Tuple[str, Union[int, str, Mapping[str, Any]]]]
-_IndexKeyHint = Union[str, _IndexList]
 
 
 class ReturnDocument(object):
@@ -1948,7 +1946,9 @@ def create_index(
     ) -> str:
         """Creates an index on this collection.
 
-        Takes either a single key or a list of (key, direction) pairs.
+        Takes either a single key or a list containing (key, direction) pairs
+        or keys.  If no direction is given, :data:`~pymongo.ASCENDING` will
+        be assumed.
         The key(s) must be an instance of :class:`basestring`
         (:class:`str` in python 3), and the direction(s) must be one of
         (:data:`~pymongo.ASCENDING`, :data:`~pymongo.DESCENDING`,
@@ -1964,7 +1964,7 @@ def create_index(
         ascending we need to use a list of tuples::
 
           >>> my_collection.create_index([("mike", pymongo.DESCENDING),
-          ...                             ("eliot", pymongo.ASCENDING)])
+          ...                             "eliot"])
 
         All optional index creation parameters should be passed as
         keyword arguments to this method. For example::
@@ -2025,6 +2025,9 @@ def create_index(
           - `**kwargs` (optional): any additional index creation
             options (see the above list) should be passed as keyword
 
+        .. versionchanged:: 4.4
+           Allow passing a list containing (key, direction) pairs
+           or keys for the ``keys`` parameter.
         .. versionchanged:: 4.1
            Added ``comment`` parameter.
         .. versionchanged:: 3.11

pymongo/cursor.py

@@ -146,7 +146,7 @@ def close(self):
             self.sock = None
 
 
-_Sort = Sequence[Tuple[str, Union[int, str, Mapping[str, Any]]]]
+_Sort = Sequence[Union[str, Tuple[str, Union[int, str, Mapping[str, Any]]]]]
 _Hint = Union[str, _Sort]
 
 
@@ -832,15 +832,16 @@ def sort(
         """Sorts this cursor's results.
 
         Pass a field name and a direction, either
-        :data:`~pymongo.ASCENDING` or :data:`~pymongo.DESCENDING`::
+        :data:`~pymongo.ASCENDING` or :data:`~pymongo.DESCENDING`.::
 
             for doc in collection.find().sort('field', pymongo.ASCENDING):
                 print(doc)
 
-        To sort by multiple fields, pass a list of (key, direction) pairs::
+        To sort by multiple fields, pass a list of (key, direction) pairs.
+        If just a name is given, :data:`~pymongo.ASCENDING` will be inferred::
 
             for doc in collection.find().sort([
-                    ('field1', pymongo.ASCENDING),
+                    'field1',
                     ('field2', pymongo.DESCENDING)]):
                 print(doc)
 

pymongo/helpers.py

@@ -80,6 +80,8 @@ def _index_list(key_or_list, direction=None):
     Takes such a list, or a single key, or a single key and direction.
     """
     if direction is not None:
+        if not isinstance(key_or_list, str):
+            raise TypeError("Expected a string and a direction")
         return [(key_or_list, direction)]
     else:
         if isinstance(key_or_list, str):
@@ -88,7 +90,12 @@ def _index_list(key_or_list, direction=None):
             return list(key_or_list)
         elif not isinstance(key_or_list, (list, tuple)):
             raise TypeError("if no direction is specified, key_or_list must be an instance of list")
-        return key_or_list
+        values = []
+        for item in key_or_list:
+            if isinstance(item, str):
+                item = (item, ASCENDING)
+            values.append(item)
+        return values
 
 
 def _index_document(index_list):
@@ -108,7 +115,10 @@ def _index_document(index_list):
         raise ValueError("key_or_list must not be the empty list")
 
     index: SON[str, Any] = SON()
-    for (key, value) in index_list:
+    for item in index_list:
+        if isinstance(item, str):
+            item = (item, ASCENDING)
+        key, value = item
         if not isinstance(key, str):
             raise TypeError("first item in each key pair must be an instance of str")
         if not isinstance(value, (str, int, abc.Mapping)):

pymongo/operations.py

@@ -22,6 +22,10 @@
 from pymongo.helpers import _gen_index_name, _index_document, _index_list
 from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
 
+# Hint supports index name, "myIndex", or list of either strings or index pairs: [('x', 1), ('y', -1), 'z'']
+_IndexList = Sequence[Union[str, Tuple[str, Union[int, str, Mapping[str, Any]]]]]
+_IndexKeyHint = Union[str, _IndexList]
+
 
 class InsertOne(Generic[_DocumentType]):
     """Represents an insert_one operation."""
@@ -55,10 +59,6 @@ def __ne__(self, other: Any) -> bool:
         return not self == other
 
 
-_IndexList = Sequence[Tuple[str, Union[int, str, Mapping[str, Any]]]]
-_IndexKeyHint = Union[str, _IndexList]
-
-
 class DeleteOne(object):
     """Represents a delete_one operation."""
 
@@ -435,7 +435,9 @@ def __init__(self, keys: _IndexKeyHint, **kwargs: Any) -> None:
 
         For use with :meth:`~pymongo.collection.Collection.create_indexes`.
 
-        Takes either a single key or a list of (key, direction) pairs.
+        Takes either a single key or a list containing (key, direction) pairs
+        or keys.  If no direction is given, :data:`~pymongo.ASCENDING` will
+        be assumed.
         The key(s) must be an instance of :class:`basestring`
         (:class:`str` in python 3), and the direction(s) must be one of
         (:data:`~pymongo.ASCENDING`, :data:`~pymongo.DESCENDING`,
@@ -477,8 +479,8 @@ def __init__(self, keys: _IndexKeyHint, **kwargs: Any) -> None:
         server version.
 
         :Parameters:
-          - `keys`: a single key or a list of (key, direction)
-            pairs specifying the index to create
+          - `keys`: a single key or a list containing (key, direction) pairs
+             or keys specifying the index to create
           - `**kwargs` (optional): any additional index creation
             options (see the above list) should be passed as keyword
             arguments

test/test_collection.py

@@ -315,6 +315,10 @@ def test_create_index(self):
         with self.write_concern_collection() as coll:
             coll.create_index([("hello", DESCENDING)])
 
+        db.test.create_index(["hello", "world"])
+        db.test.create_index(["hello", ("world", DESCENDING)])
+        db.test.create_index({"hello": 1}.items())  # type:ignore[arg-type]
+
     def test_drop_index(self):
         db = self.db
         db.test.drop_indexes()
@@ -1680,7 +1684,7 @@ def to_list(things):
 
         self.assertRaises(TypeError, db.test.find, sort=5)
         self.assertRaises(TypeError, db.test.find, sort="hello")
-        self.assertRaises(ValueError, db.test.find, sort=["hello", 1])
+        self.assertRaises(TypeError, db.test.find, sort=["hello", 1])
 
     # TODO doesn't actually test functionality, just that it doesn't blow up
     def test_cursor_timeout(self):

test/test_cursor.py

@@ -40,6 +40,7 @@
 from pymongo.collation import Collation
 from pymongo.cursor import Cursor, CursorType
 from pymongo.errors import ExecutionTimeout, InvalidOperation, OperationFailure
+from pymongo.operations import _IndexList
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
@@ -366,6 +367,21 @@ def test_hint(self):
             break
         self.assertRaises(InvalidOperation, a.hint, spec)
 
+        db.test.drop()
+        db.test.insert_many([{"num": i, "foo": i} for i in range(100)])
+        spec: _IndexList = ["num", ("foo", DESCENDING)]
+        db.test.create_index(spec)
+        first = next(db.test.find().hint(spec))
+        self.assertEqual(0, first.get("num"))
+        self.assertEqual(0, first.get("foo"))
+
+        db.test.drop()
+        db.test.insert_many([{"num": i, "foo": i} for i in range(100)])
+        spec = ["num"]
+        db.test.create_index(spec)
+        first = next(db.test.find().hint(spec))
+        self.assertEqual(0, first.get("num"))
+
     def test_hint_by_name(self):
         db = self.db
         db.test.drop()
@@ -715,6 +731,8 @@ def test_sort(self):
             (i["a"], i["b"]) for i in db.test.find().sort([("b", DESCENDING), ("a", ASCENDING)])
         ]
         self.assertEqual(result, expected)
+        result = [(i["a"], i["b"]) for i in db.test.find().sort([("b", DESCENDING), "a"])]
+        self.assertEqual(result, expected)
 
         a = db.test.find()
         a.sort("x", ASCENDING)