Add spec test functionality

.travis.yml

@@ -3,6 +3,8 @@ language: python
 services:
   - mongodb
 
+install: python -m pip install https://github.com/mongodb/mongo-python-driver/archive/3.11.0rc0.tar.gz
+
 python:
   - 3.5
   - 3.6

pymongoexplain/commands.py

@@ -19,16 +19,71 @@
 from typing import Union
 
 from bson.son import SON
+from bson.py3compat import abc, iteritems, string_type
+
 from pymongo.collection import Collection
+from pymongo.collation import validate_collation_or_none
 from .utils import convert_to_camelcase
 
 
 Document = Union[dict, SON]
 
 
+def _index_document(index_list):
+    """Helper to generate an index specifying document.
+
+    Takes a list of (key, direction) pairs.
+    """
+    if isinstance(index_list, abc.Mapping):
+        raise TypeError("passing a dict to sort/create_index/hint is not "
+                        "allowed - use a list of tuples instead. did you "
+                        "mean %r?" % list(iteritems(index_list)))
+    elif not isinstance(index_list, (list, tuple)):
+        raise TypeError("must use a list of (key, direction) pairs, "
+                        "not: " + repr(index_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, string_type):
+            raise TypeError("first item in each key pair must be a string")
+        if not isinstance(value, (string_type, int, abc.Mapping)):
+            raise TypeError("second item in each key pair must be 1, -1, "
+                            "'2d', 'geoHaystack', or another valid MongoDB "
+                            "index specifier.")
+        index[key] = value
+    return index
+
+
+def _fields_list_to_dict(fields, option_name):
+    """Takes a sequence of field names and returns a matching dictionary.
+
+    ["a", "b"] becomes {"a": 1, "b": 1}
+
+    and
+
+    ["a.b.c", "d", "a.c"] becomes {"a.b.c": 1, "d": 1, "a.c": 1}
+    """
+    if isinstance(fields, abc.Mapping):
+        return fields
+
+    if isinstance(fields, (abc.Sequence, abc.Set)):
+        if not all(isinstance(field, string_type) for field in fields):
+            raise TypeError("%s must be a list of key names, each an "
+                            "instance of %s" % (option_name,
+                                                string_type.__name__))
+        return dict.fromkeys(fields, 1)
+
+    raise TypeError("%s must be a mapping or "
+                    "list of key names" % (option_name,))
+
 class BaseCommand():
-    def __init__(self, collection):
+    def __init__(self, collection, collation):
         self.command_document = {}
+        collation = validate_collation_or_none(collation)
+        if collation is not None:
+            self.command_document["collation"] = collation
         self.collection = collection
 
     @property
@@ -44,29 +99,51 @@ def get_SON(self):
 
 class UpdateCommand(BaseCommand):
     def __init__(self, collection: Collection, filter, update,
-                 kwargs):
-        super().__init__(collection.name)
-        return_document = {"updates":[{"q": filter, "u": update}]}
-        for key in kwargs:
-            value = kwargs[key]
-            if key == "bypass_document_validation":
-                return_document[key] = value
-            else:
-                return_document["updates"][0][key] = value
-        self.command_document = convert_to_camelcase(return_document)
+                 upsert=None, multi=None, collation=None, array_filters=None,
+                 hint=None, ordered=None, write_concern=None,
+                 bypass_document_validation=None, comment=None):
+        super().__init__(collection.name, collation)
+        update_doc = {"q": filter, "u": update}
+        if upsert is not None:
+            update_doc["upsert"] = upsert
+
+        if multi is not None:
+            update_doc["multi"] = multi
+
+        if array_filters is not None:
+            update_doc["array_filters"] = array_filters
+
+        if hint is not None:
+            update_doc["hint"] = hint if \
+                        isinstance(hint, str) else _index_document(hint)
+        self.command_document["updates"] = [update_doc]
+
+        if ordered is not None:
+            self.command_document["ordered"] = ordered
+
+        if write_concern is not None:
+            self.command_document["write_concern"] = write_concern
+
+        if bypass_document_validation is not None and \
+                bypass_document_validation is not False:
+            self.command_document["bypass_document_validation"] = bypass_document_validation
+
+        if comment is not None:
+            self.command_document["comment"] = comment
+
+        self.command_document = convert_to_camelcase(self.command_document)
 
     @property
     def command_name(self):
         return "update"
 
 
 class DistinctCommand(BaseCommand):
-    def __init__(self, collection: Collection, key, filter, session,
+    def __init__(self, collection: Collection, key, filter,
                  kwargs):
-        super().__init__(collection.name)
-        self.command_document = {"key": key, "query": filter}
-        for key, value in kwargs.items():
-            self.command_document[key] = value
+        super().__init__(collection.name, kwargs.pop("collation", None))
+        self.command_document.update({"key": key, "query": filter})
+
         self.command_document = convert_to_camelcase(self.command_document)
 
     @property
@@ -75,26 +152,34 @@ def command_name(self):
 
 
 class AggregateCommand(BaseCommand):
-    def __init__(self, collection: Collection, pipeline, session,
+    def __init__(self, collection: Collection, pipeline,
                  cursor_options,
-                 kwargs, exclude_keys = []):
-        super().__init__(collection.name)
-        self.command_document = {"pipeline": pipeline, "cursor": cursor_options}
+                 kwargs):
+
+        super().__init__(collection.name, kwargs.pop("collation", None))
+        self.command_document.update({"pipeline": pipeline, "cursor":
+            cursor_options})
+
         for key, value in kwargs.items():
-            self.command_document[key] = value
+            if key == "batchSize":
+                if value == 0:
+                    continue
+                self.command_document["cursor"]["batchSize"] = value
+            else:
+                self.command_document[key] = value
 
         self.command_document = convert_to_camelcase(
-            self.command_document, exclude_keys=exclude_keys)
+            self.command_document)
 
     @property
     def command_name(self):
         return "aggregate"
 
+
 class CountCommand(BaseCommand):
-    def __init__(self, collection: Collection, filter,
-                 kwargs):
-        super().__init__(collection.name)
-        self.command_document = {"query": filter}
+    def __init__(self, collection: Collection, filter, kwargs):
+        super().__init__(collection.name, kwargs.pop("collation", None))
+        self.command_document.update({"query": filter})
         for key, value in kwargs.items():
             self.command_document[key] = value
         self.command_document = convert_to_camelcase(self.command_document)
@@ -107,21 +192,37 @@ def command_name(self):
 class FindCommand(BaseCommand):
     def __init__(self, collection: Collection,
                  kwargs):
-        super().__init__(collection.name)
+        super().__init__(collection.name, kwargs.pop("collation", None))
         for key, value in kwargs.items():
-            self.command_document[key] = value
+            if key == "projection" and value is not None:
+                self.command_document["projection"] = _fields_list_to_dict(
+                    value, "projection")
+            elif key == "sort":
+                self.command_document["sort"] = _index_document(
+                    value)
+            else:
+                self.command_document[key] = value
+
         self.command_document = convert_to_camelcase(self.command_document)
 
     @property
     def command_name(self):
         return "find"
 
+
 class FindAndModifyCommand(BaseCommand):
     def __init__(self, collection: Collection,
                  kwargs):
-        super().__init__(collection.name)
+        super().__init__(collection.name, kwargs.pop("collation", None))
         for key, value in kwargs.items():
-            self.command_document[key] = value
+            if key == "hint":
+                self.command_document["hint"] = value if \
+                    isinstance(value, str) else _index_document(value)
+            elif key == "sort" and value is not None:
+                self.command_document["sort"] = _index_document(
+                    value)
+            else:
+                self.command_document[key] = value
         self.command_document = convert_to_camelcase(self.command_document)
 
     @property
@@ -132,10 +233,16 @@ def command_name(self):
 class DeleteCommand(BaseCommand):
     def __init__(self, collection: Collection, filter,
                  limit, collation, kwargs):
-        super().__init__(collection.name)
-        self.command_document = {"deletes": [SON({"q": filter, "limit": limit})]}
+        super().__init__(collection.name, kwargs.pop("collation", None))
+        self.command_document["deletes"] = [{"q": filter, "limit":
+            limit}]
         for key, value in kwargs.items():
-            self.command_document[key] = value
+            if key == "hint":
+                self.command_document["deletes"][0]["hint"] = value if \
+                    isinstance(value, str) else _index_document(value)
+            else:
+                self.command_document[key] = value
+
         self.command_document = convert_to_camelcase(self.command_document)
 
     @property