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

@@ -20,15 +20,20 @@
 
 from bson.son import SON
 from pymongo.collection import Collection
+from pymongo.helpers import _index_document, _fields_list_to_dict
+from pymongo.collation import validate_collation_or_none
 from .utils import convert_to_camelcase
 
 
 Document = Union[dict, SON]
 
 
 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 +49,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}
+        self.command_document["updates"] = [update_doc]
+        if upsert is not None:
+            self.command_document["updates"][0]["upsert"] = upsert
+
+        if multi is not None:
+            self.command_document["updates"][0]["multi"] = multi
+
+        if array_filters is not None:
+            self.command_document["updates"][0]["array_filters"] = array_filters
+
+        if hint is not None:
+            self.command_document["updates"][0]["hint"] = hint if \
+                        isinstance(hint, str) else _index_document(hint)
+
+        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 +102,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 +142,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 +183,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

pymongoexplain/explainable_collection.py

@@ -16,13 +16,15 @@
 from typing import Union, List, Dict
 
 import pymongo
+from pymongo.collection import Collection
 from bson.son import SON
 
 from .commands import AggregateCommand, FindCommand, CountCommand, \
     UpdateCommand, DistinctCommand, DeleteCommand, FindAndModifyCommand
 
 Document = Union[dict, SON]
 
+
 class ExplainCollection():
     def __init__(self, collection):
         self.collection = collection
@@ -39,31 +41,28 @@ def update_one(self, filter, update, upsert=False,
                    bypass_document_validation=False,
                    collation=None, array_filters=None, hint=None,
                    session=None, **kwargs):
-        kwargs.update(locals())
-        del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs["update"]
-        kwargs["multi"] = False
-        if bypass_document_validation == False:
-            del kwargs["bypass_document_validation"]
-        command = UpdateCommand(self.collection, filter, update, kwargs)
+        command = UpdateCommand(self.collection, filter, update,
+                                bypass_document_validation=
+                                    bypass_document_validation,
+                                array_filters=array_filters,
+                                collation=collation, hint=hint,
+                                upsert=upsert, multi=False)
         return self._explain_command(command)
 
     def update_many(self, filter: Document, update: Document, upsert=False,
-                    array_filters=None, bypass_document_validation=False, collation=None, session=None, **kwargs):
-        kwargs.update(locals())
-        del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs["update"]
-        kwargs["multi"] = True
-        if bypass_document_validation == False:
-            del kwargs["bypass_document_validation"]
-        command = UpdateCommand(self.collection, filter, update, kwargs)
+                    array_filters=None, bypass_document_validation=False,
+                    collation=None, hint=None, session=None, **kwargs):
+        command = UpdateCommand(self.collection, filter, update, multi=True,
+                                bypass_document_validation=bypass_document_validation, upsert=upsert, collation=collation, array_filters=array_filters, hint=hint)
         return self._explain_command(command)
 
     def distinct(self, key: str, filter: Document=None, session=None, **kwargs):
-        command = DistinctCommand(self.collection, key, filter, session, kwargs)
+        command = DistinctCommand(self.collection, key, filter, kwargs)
         return self._explain_command(command)
 
     def aggregate(self, pipeline: List[Document], session=None, **kwargs):
-        command = AggregateCommand(self.collection, pipeline, session,
-                                   {},kwargs)
+        command = AggregateCommand(self.collection, pipeline,
+                                   {}, kwargs)
         return self._explain_command(command)
 
     def estimated_document_count(self,
@@ -76,9 +75,8 @@ def count_documents(self, filter: Document, session=None,
                                  **kwargs):
 
         command = AggregateCommand(self.collection, [{'$match': filter},
-                                                      {'$group': {'n': {'$sum': 1}, '_id': 1}}],
-                                   session, {}, kwargs,
-                                   exclude_keys=filter.keys())
+                                                     {'$group': {'n': {'$sum': 1}, '_id': 1}}]
+                                   , {}, kwargs)
         return self._explain_command(command)
 
     def delete_one(self, filter: Document, collation=None, session=None,
@@ -93,16 +91,14 @@ def delete_many(self, filter: Document, collation=None,
                                                             Document,
                                                       bool]]):
         limit = 0
-        kwargs["session"] = session
         command = DeleteCommand(self.collection, filter, limit, collation,
-        kwargs)
+                                kwargs)
         return self._explain_command(command)
 
     def watch(self, pipeline: Document = None, full_document: Document = None,
               resume_after= None,
               max_await_time_ms: int = None, batch_size: int = None,
-              collation=None, start_at_operation_time=None, session:
-            pymongo.mongo_client.client_session.ClientSession=None,
+              collation=None, start_at_operation_time=None, session=None,
               start_after=None):
         change_stream_options = {"start_after":start_after,
                                  "resume_after":resume_after,
@@ -114,7 +110,7 @@ def watch(self, pipeline: Document = None, full_document: Document = None,
             pipeline = [{"$changeStream": change_stream_options}]
 
         command = AggregateCommand(self.collection, pipeline,
-                                   session, {"batch_size":batch_size},
+                                    {"batch_size":batch_size},
                                    {"collation":collation, "max_await_time_ms":
                                        max_await_time_ms})
         return self._explain_command(command)
@@ -124,7 +120,8 @@ def find(self, filter: Document = None,
         kwargs.update(locals())
         del kwargs["self"], kwargs["kwargs"]
         command = FindCommand(self.collection,
-                                kwargs)
+                              kwargs)
+
         return self._explain_command(command)
 
     def find_one(self, filter: Document = None, **kwargs: Dict[str,
@@ -149,7 +146,8 @@ def find_one_and_delete(self, filter: Document, projection: list = None,
                                        kwargs)
         return self._explain_command(command)
 
-    def find_one_and_replace(self, filter: Document, replacement: Document,
+    def find_one_and_replace(self, filter: Document, replacement:
+    Document={},
                             projection: list = None, sort=None,
                              return_document=pymongo.ReturnDocument.BEFORE,
                              session=None, **kwargs):
@@ -163,15 +161,15 @@ def find_one_and_replace(self, filter: Document, replacement: Document,
                                        kwargs)
         return self._explain_command(command)
 
-    def find_one_and_update(self, filter: Document, replacement: Document,
+    def find_one_and_update(self, filter: Document, update: Document,
                             projection: list = None, sort=None,
                              return_document=pymongo.ReturnDocument.BEFORE,
                              session=None, **kwargs):
         kwargs["query"] = filter
         kwargs["fields"] = projection
         kwargs["sort"] = sort
         kwargs["upsert"] = False
-        kwargs["update"] = replacement
+        kwargs["update"] = update
         kwargs["session"] = session
 
         command = FindAndModifyCommand(self.collection,
@@ -180,15 +178,10 @@ def find_one_and_update(self, filter: Document, replacement: Document,
 
     def replace_one(self, filter: Document, replacement: Document,
                     upsert=False, bypass_document_validation=False,
-                    collation=None, session=None, **kwargs):
-        kwargs.update(locals())
-        del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs[
-            "replacement"]
-        kwargs["multi"] = False
-        if not bypass_document_validation:
-            del kwargs["bypass_document_validation"]
-        update = replacement
-        command = UpdateCommand(self.collection, filter, update, kwargs)
+                    collation=None, hint=None, session=None, **kwargs):
+        command = UpdateCommand(self.collection, filter, update=replacement,
+                                bypass_document_validation=bypass_document_validation,
+                                hint=hint, collation=collation, multi=False, upsert=upsert)
 
         return self._explain_command(command)
 

pymongoexplain/utils.py

@@ -16,26 +16,17 @@
 """Utility functions"""
 
 
-def convert_to_camelcase(d, exclude_keys=[]):
+def convert_to_camelcase(d):
     if not isinstance(d, dict):
         return d
     ret = dict()
     for key in d.keys():
         if d[key] is None:
             continue
-        if key in exclude_keys:
-            ret[key] = d[key]
-            continue
         new_key = key
         if "_" in key and key[0] != "_":
             new_key = key.split("_")[0] + ''.join(
                 [i.capitalize() for i in key.split("_")[1:]])
-        if isinstance(d[key], list):
-            ret[new_key] = [convert_to_camelcase(
-                i, exclude_keys=exclude_keys) for i in d[key]]
-        elif isinstance(d[key], dict):
-            ret[new_key] = convert_to_camelcase(d[key],
-                                                     exclude_keys=exclude_keys)
         else:
             ret[new_key] = d[key]
     return ret