"""Definition of function to count documents in database."""

from __future__ import annotations

from pymongo.collection import Collection


def count_documents(
    collection: Collection,
    only_with_annotation: bool = False,
) -> int:
    """Get the total number of documents in database that matches the
    filters."""
    assert isinstance(collection, Collection)
    assert isinstance(only_with_annotation, bool)
    # build query
    query = {}
    if only_with_annotation:
        query['annotation'] = {'$ne': None}
    return collection.count_documents(filter=query)