From 742d19e2204610a436283c3b73781258e390cc43 Mon Sep 17 00:00:00 2001 From: Dillon Walls Date: Fri, 28 Jun 2024 15:55:30 -0400 Subject: [PATCH] drop mongo's count() methods, but keep ming session and cursor count methods --- ming/base.py | 5 +- ming/mim.py | 24 ++++---- ming/odm/odmsession.py | 1 + ming/session.py | 5 +- ming/tests/test_declarative.py | 3 +- ming/tests/test_functional.py | 3 +- ming/tests/test_mim.py | 109 ++++++++++++++++----------------- 7 files changed, 79 insertions(+), 71 deletions(-) diff --git a/ming/base.py b/ming/base.py index aa6229f..de3073f 100644 --- a/ming/base.py +++ b/ming/base.py @@ -69,11 +69,12 @@ class Cursor: def __bool__(self): raise MingException('Cannot evaluate Cursor to a boolean') - def __init__(self, cls, cursor, allow_extra=True, strip_extra=True): + def __init__(self, cls, cursor, allow_extra=True, strip_extra=True, find_spec=None): self.cls = cls self.cursor = cursor self._allow_extra = allow_extra self._strip_extra = strip_extra + self.find_spec = find_spec def __iter__(self): return self @@ -89,7 +90,7 @@ def next(self): __next__ = next def count(self): - return self.cursor.count() + return self.cursor.collection.count_documents(self.find_spec) def distinct(self, *args, **kwargs): return self.cursor.distinct(*args, **kwargs) diff --git a/ming/mim.py b/ming/mim.py index 44ac0da..d47b37a 100644 --- a/ming/mim.py +++ b/ming/mim.py @@ -461,8 +461,11 @@ def find_one_and_update(self, filter, update, projection=None, sort=None, upsert sort=sort, upsert=upsert, return_document=return_document, operation=ModifyOperation.UPDATE, **kwargs) - def count(self, filter=None, **kwargs): - return self.find(filter, **kwargs).count() + def estimated_document_count(self, **kwargs): + return self.find({}, **kwargs)._count() + + def count_documents(self, filter=None, **kwargs): + return self.find(filter, **kwargs)._count() def __insert(self, doc_or_docs, **kwargs): result = [] @@ -681,7 +684,7 @@ def __init__(self, collection, _iterator_gen, if isinstance(projection, (tuple, list)): projection = {f: 1 for f in projection} - self._collection = collection + self.collection = collection self._iterator_gen = _iterator_gen self._sort = sort self._skip = skip or None # cope with 0 being passed. @@ -706,7 +709,7 @@ def iterator(self): def clone(self, **overrides): result = Cursor( - collection=self._collection, + collection=self.collection, _iterator_gen=self._iterator_gen, sort=self._sort, skip=self._skip, @@ -722,10 +725,9 @@ def rewind(self): del self.iterator self._safe_to_chain = True + def _count(self): """ - replace uses with collection based checks: - ntotal = collection.estimated_document_count() - nmatched = collection.count_documents({'price': {'$gte': 10}}) + An internal method to count the number of documents in the cursor. """ return sum(1 for x in self._iterator_gen()) @@ -756,7 +758,7 @@ def next(self): # mim doesn't currently do anything with codec_options, so this doesn't do anything currently # but leaving it here as a placeholder for the future - otherwise we should delete wrap_as_class() - return wrap_as_class(value, self._collection.codec_options.document_class) + return wrap_as_class(value, self.collection.codec_options.document_class) __next__ = next @@ -803,13 +805,13 @@ def hint(self, index): # checks indexes, but doesn't actually use hinting if type(index) == list: test_idx = [(i, direction) for i, direction in index if i != '$natural'] - values = [[k for k in i["key"]] for i in self._collection._indexes.values()] + values = [[k for k in i["key"]] for i in self.collection._indexes.values()] if test_idx and test_idx not in values: raise OperationFailure('database error: bad hint. Valid values: %s' % values) elif isinstance(index, str): - if index not in self._collection._indexes.keys(): + if index not in self.collection._indexes.keys(): raise OperationFailure('database error: bad hint. Valid values: %s' - % self._collection._indexes.keys()) + % self.collection._indexes.keys()) elif index is None: pass else: diff --git a/ming/odm/odmsession.py b/ming/odm/odmsession.py index d1f814f..2b4cb87 100644 --- a/ming/odm/odmsession.py +++ b/ming/odm/odmsession.py @@ -1,4 +1,5 @@ from collections import defaultdict +import warnings from pymongo.collection import ReturnDocument from pymongo.database import Database diff --git a/ming/session.py b/ming/session.py index 418e715..7f1a656 100644 --- a/ming/session.py +++ b/ming/session.py @@ -83,12 +83,15 @@ def find(self, cls, *args, **kwargs): collection = self._impl(cls) cursor = collection.find(*args, **kwargs) + find_spec = kwargs.get('filter', None) or args[0] if args else {} + if not validate: return (cls(o, skip_from_bson=True) for o in cursor) return Cursor(cls, cursor, allow_extra=allow_extra, - strip_extra=strip_extra) + strip_extra=strip_extra, + find_spec=find_spec) def remove(self, cls, filter={}, *args, **kwargs): fix_write_concern(kwargs) diff --git a/ming/tests/test_declarative.py b/ming/tests/test_declarative.py index 9665418..1328b05 100644 --- a/ming/tests/test_declarative.py +++ b/ming/tests/test_declarative.py @@ -330,7 +330,8 @@ class __mongometa__: b=Field(S.Object(dict(a=int))) self.TestDoc = TestDoc mongo_cursor = IteratorMock(iter([ {}, {}, {} ])) - mongo_cursor.count = mock.Mock(return_value=3) + mongo_cursor.collection = mock.Mock() + mongo_cursor.collection.count_documents = mock.Mock(return_value=3) mongo_cursor.limit = mock.Mock(return_value=mongo_cursor) mongo_cursor.hint = mock.Mock(return_value=mongo_cursor) mongo_cursor.skip = mock.Mock(return_value=mongo_cursor) diff --git a/ming/tests/test_functional.py b/ming/tests/test_functional.py index 56e8603..9cb7aa9 100644 --- a/ming/tests/test_functional.py +++ b/ming/tests/test_functional.py @@ -201,7 +201,8 @@ def next(self): Field('b', dict(a=int))) mongo_cursor = IteratorMock(iter([ {}, {}, {} ])) - mongo_cursor.count = mock.Mock(return_value=3) + mongo_cursor.collection = mock.Mock() + mongo_cursor.collection.count_documents = mock.Mock(return_value=3) mongo_cursor.limit = mock.Mock(return_value=mongo_cursor) mongo_cursor.hint = mock.Mock(return_value=mongo_cursor) mongo_cursor.skip = mock.Mock(return_value=mongo_cursor) diff --git a/ming/tests/test_mim.py b/ming/tests/test_mim.py index 9811358..42c2b65 100644 --- a/ming/tests/test_mim.py +++ b/ming/tests/test_mim.py @@ -27,79 +27,79 @@ def test_limit(self): self.assertEqual(4, len(f({}).limit(0).all())) def test_regex(self): - f = self.bind.db.rcoll.find - assert 4 == f(dict(_id=re.compile(r'r\d+'))).count() - assert 2 == f(dict(_id=re.compile(r'r[0-1]'))).count() + f = self.bind.db.rcoll.count_documents + assert 4 == f(dict(_id=re.compile(r'r\d+'))) + assert 2 == f(dict(_id=re.compile(r'r[0-1]'))) def test_regex_options(self): - f = self.bind.db.rcoll.find - assert 2 == f(dict(_id={'$regex': 'r[0-1]', '$options': 'i'})).count() + f = self.bind.db.rcoll.count_documents + assert 2 == f(dict(_id={'$regex': 'r[0-1]', '$options': 'i'})) def test_eq(self): - f = self.bind.db.rcoll.find - assert 1 == f(dict(d={'$eq': 0})).count() + f = self.bind.db.rcoll.count_documents + assert 1 == f(dict(d={'$eq': 0})) def test_ne(self): - f = self.bind.db.rcoll.find - assert 3 == f(dict(d={'$ne': 0})).count() + f = self.bind.db.rcoll.count_documents + assert 3 == f(dict(d={'$ne': 0})) def test_gt(self): - f = self.bind.db.rcoll.find - assert 1 == f(dict(d={'$gt': 2})).count() - assert 0 == f(dict(d={'$gt': 3})).count() + f = self.bind.db.rcoll.count_documents + assert 1 == f(dict(d={'$gt': 2})) + assert 0 == f(dict(d={'$gt': 3})) def test_gte(self): - f = self.bind.db.rcoll.find - assert 2 == f(dict(d={'$gte': 2})).count() - assert 1 == f(dict(d={'$gte': 3})).count() + f = self.bind.db.rcoll.count_documents + assert 2 == f(dict(d={'$gte': 2})) + assert 1 == f(dict(d={'$gte': 3})) def test_lt(self): - f = self.bind.db.rcoll.find - assert 0 == f(dict(d={'$lt': 0})).count() - assert 1 == f(dict(d={'$lt': 1})).count() - assert 2 == f(dict(d={'$lt': 2})).count() + f = self.bind.db.rcoll.count_documents + assert 0 == f(dict(d={'$lt': 0})) + assert 1 == f(dict(d={'$lt': 1})) + assert 2 == f(dict(d={'$lt': 2})) def test_lte(self): - f = self.bind.db.rcoll.find - assert 1 == f(dict(d={'$lte': 0})).count() - assert 2 == f(dict(d={'$lte': 1})).count() - assert 3 == f(dict(d={'$lte': 2})).count() + f = self.bind.db.rcoll.count_documents + assert 1 == f(dict(d={'$lte': 0})) + assert 2 == f(dict(d={'$lte': 1})) + assert 3 == f(dict(d={'$lte': 2})) def test_range_equal(self): - f = self.bind.db.rcoll.find - assert 1 == f(dict(d={'$gte': 2, '$lte': 2})).count() - assert 2 == f(dict(d={'$gte': 1, '$lte': 2})).count() - assert 0 == f(dict(d={'$gte': 4, '$lte': -1})).count() + f = self.bind.db.rcoll.count_documents + assert 1 == f(dict(d={'$gte': 2, '$lte': 2})) + assert 2 == f(dict(d={'$gte': 1, '$lte': 2})) + assert 0 == f(dict(d={'$gte': 4, '$lte': -1})) def test_range_inequal(self): - f = self.bind.db.rcoll.find - assert 0 == f(dict(d={'$gt': 2, '$lt': 2})).count() - assert 1 == f(dict(d={'$gt': 2, '$lt': 4})).count() - assert 0 == f(dict(d={'$gt': 1, '$lt': 2})).count() - assert 1 == f(dict(d={'$gt': 1, '$lt': 3})).count() - assert 0 == f(dict(d={'$gt': 4, '$lt': -1})).count() + f = self.bind.db.rcoll.count_documents + assert 0 == f(dict(d={'$gt': 2, '$lt': 2})) + assert 1 == f(dict(d={'$gt': 2, '$lt': 4})) + assert 0 == f(dict(d={'$gt': 1, '$lt': 2})) + assert 1 == f(dict(d={'$gt': 1, '$lt': 3})) + assert 0 == f(dict(d={'$gt': 4, '$lt': -1})) def test_exists(self): - f = self.bind.db.coll.find - assert 1 == f(dict(a={'$exists':True})).count() - assert 0 == f(dict(a={'$exists':False})).count() - assert 0 == f(dict(b={'$exists':True})).count() - assert 1 == f(dict(b={'$exists':False})).count() + f = self.bind.db.coll.count_documents + assert 1 == f(dict(a={'$exists':True})) + assert 0 == f(dict(a={'$exists':False})) + assert 0 == f(dict(b={'$exists':True})) + assert 1 == f(dict(b={'$exists':False})) def test_all(self): - f = self.bind.db.coll.find - assert 1 == f(dict(c={'$all':[1,2]})).count() - assert 1 == f(dict(c={'$all':[1,2,3]})).count() - assert 0 == f(dict(c={'$all':[2,3,4]})).count() - assert 1 == f(dict(c={'$all':[]})).count() + f = self.bind.db.coll.count_documents + assert 1 == f(dict(c={'$all':[1,2]})) + assert 1 == f(dict(c={'$all':[1,2,3]})) + assert 0 == f(dict(c={'$all':[2,3,4]})) + assert 1 == f(dict(c={'$all':[]})) def test_or(self): - f = self.bind.db.coll.find - assert 1 == f(dict({'$or': [{'c':{'$all':[1,2,3]}}]})).count() - assert 0 == f(dict({'$or': [{'c':{'$all':[4,2,3]}}]})).count() - assert 1 == f(dict({'$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]})).count() - self.assertEqual(0, f(dict({'_id': 'bar', '$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]})).count()) - self.assertEqual(1, f(dict({'_id': 'foo', '$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]})).count()) + f = self.bind.db.coll.count_documents + assert 1 == f(dict({'$or': [{'c':{'$all':[1,2,3]}}]})) + assert 0 == f(dict({'$or': [{'c':{'$all':[4,2,3]}}]})) + assert 1 == f(dict({'$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]})) + self.assertEqual(0, f(dict({'_id': 'bar', '$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]}))) + self.assertEqual(1, f(dict({'_id': 'foo', '$or': [{'a': 2}, {'c':{'$all':[1,2,3]}}]}))) def test_find_with_projection_list(self): o = self.bind.db.coll.find_one({'a': 2}, projection=['a']) @@ -187,8 +187,7 @@ def test_search(self): coll.create_index([('field', 'text')]) coll.insert_one({'field': 'text to be searched'}) coll.insert_one({'field': 'text to be'}) - assert coll.find({'$text': {'$search': 'searched'}}, - {'score': {'$meta': 'textScore'}}).count() == 1 + assert coll.count_documents({'$text': {'$search': 'searched'}}) == 1 class TestDottedOperators(TestCase): @@ -215,8 +214,8 @@ def test_inc_dotted_dollar_middle1(self): self.assertEqual(obj, { 'b': { 'f': [ { 'g': 11 }, { 'g': 2 } ] }}) def test_find_dotted(self): - self.assertEqual(self.coll.find({'b.c': 1}).count(), 1) - self.assertEqual(self.coll.find({'b.c': 2}).count(), 0) + self.assertEqual(self.coll.count_documents({'b.c': 1}), 1) + self.assertEqual(self.coll.count_documents({'b.c': 2}), 0) self.assertEqual(0, len(self.coll.find({'x.y.z': 1}).all())) def test_inc_dotted(self): @@ -596,13 +595,13 @@ def test_find_one_and_replace_returns_new_value_on_new_upsert(self): def test_find_one_and_delete(self): self.bind.db.col.insert_one({'_id': 1}) self.assertEqual({'_id': 1}, self.bind.db.col.find_one_and_delete({'_id': 1})) - self.assertEqual(0, self.bind.db.col.count()) + self.assertEqual(0, self.bind.db.col.estimated_document_count()) def test_find_one_and_delete_returns_projection(self): self.bind.db.col.insert_one({'_id': 1, 'i': 1}) self.assertEqual({'i': 1}, self.bind.db.col.find_one_and_delete({'_id': 1}, projection={'_id': False, 'i': True})) - self.assertEqual(0, self.bind.db.col.count()) + self.assertEqual(0, self.bind.db.col.estimated_document_count()) def test_hint_simple(self): self.bind.db.coll.ensure_index([('myindex', 1)])