letsswing_admin/.venv/lib/python3.11/site-packages/montydb/collection.py

import warnings
from copy import deepcopy

from .base import (
    BaseObject,
    WriteConcern,
    validate_is_mapping,
    validate_ok_for_update,
    validate_ok_for_replace,
    validate_list_or_none,
    validate_boolean,
)

from .cursor import MontyCursor
from .engine.field_walker import FieldWalker
from .engine.weighted import Weighted
from .engine.queries import QueryFilter
from .engine.update import Updator
from .types import (
    abc,
    bson,
    string_types,
    is_duckument_type,
    Counter,
    on_err_close,
)
from .storage import StorageDuplicateKeyError
from .errors import (
    DuplicateKeyError,
    BulkWriteError,
    WriteError,
)

from .results import (
    DeleteResult,
    InsertOneResult,
    InsertManyResult,
    UpdateResult,
)


NotImplementeds = {
    "aggregate",
    "aggregate_raw_batches",
    "bulk_write",
    "watch",
    "find_raw_batches",
    "find_one_and_delete",
    "find_one_and_replace",
    "find_one_and_update",
    "create_indexes",
    "drop_index",
    "drop_indexes",
    "reindex",
    "list_indexes",
    "index_information",
    "rename",
    "options",
    "map_reduce",
    "inline_map_reduce",
    "parallel_scan",
}


class MontyCollection(BaseObject):
    def __init__(
        self,
        database,
        name,
        create=False,
        codec_options=None,
        read_preference=None,
        write_concern=None,
        read_concern=None,
        session=None,
        **kwargs
    ):
        """ """
        super(MontyCollection, self).__init__(
            codec_options or database.codec_options,
            write_concern or database.write_concern,
        )

        self._storage = database.client._storage

        self._database = database
        self._name = name
        self._components = (database, self)

    def __repr__(self):
        return "MontyCollection({!r}, {!r})".format(self._database, self._name)

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return self._database == other.database and self._name == other.name
        return NotImplemented

    def __ne__(self, other):
        return not self == other

    def __getattr__(self, name):
        if name in NotImplementeds:
            raise NotImplementedError(
                "'MontyCollection.%s' is NOT implemented !" % name
            )
        if name.startswith("_"):
            full_name = ".".join((self._name, name))
            raise AttributeError(
                "MontyCollection has no attribute {0!r}. To access the {1}"
                " collection, use database[{1!r}].".format(name, full_name)
            )
        return self.__getitem__(name)

    def __getitem__(self, key):
        return self._database.get_collection(".".join((self._name, key)))

    @property
    def full_name(self):
        """ """
        return u".".join((self._database.name, self._name))

    @property
    def name(self):
        """ """
        return self._name

    @property
    def database(self):
        """ """
        return self._database

    def with_options(self, codec_options=None, write_concern=None, *args, **kwargs):
        if not isinstance(write_concern, WriteConcern):
            # Could be `pymongo.WriteConcern` if called from mongoengine.
            write_concern = None

        return MontyCollection(
            self._database,
            self._name,
            False,
            codec_options or self.codec_options,
            write_concern or self.write_concern,
        )

    def insert_one(self, document, bypass_document_validation=False, *args, **kwargs):
        """ """
        if bypass_document_validation:
            pass

        if "_id" not in document:
            document["_id"] = bson.ObjectId()

        try:
            result = self._storage.write_one(self, document)
        except StorageDuplicateKeyError:
            message = (
                "E11000 duplicate key error collection: %s index: "
                '_id_ dup key: { : "%s" }' % (self.full_name, str(document["_id"]))
            )
            details = {"index": 0, "code": 11000, "errmsg": message}
            raise DuplicateKeyError(message, code=11000, details=details)

        return InsertOneResult(result)

    def insert_many(
        self, documents, ordered=True, bypass_document_validation=False, *args, **kwargs
    ):
        """ """
        if not isinstance(documents, abc.Iterable) or not documents:
            raise TypeError("documents must be a non-empty list")

        if bypass_document_validation:
            pass

        def set_id(doc):
            if "_id" not in doc:
                doc["_id"] = bson.ObjectId()
            # Keep _id in track for error message
            return doc["_id"]

        counter = Counter(iter(documents), job_on_each=set_id)

        try:
            result = self._storage.write_many(self, counter, ordered)
        except StorageDuplicateKeyError:
            message = (
                "E11000 duplicate key error collection: %s index: "
                '_id_ dup key: { : "%s" }' % (self.full_name, str(counter.data))
            )
            index = counter.count - 1
            result = {
                "writeErrors": [
                    {
                        "index": index,
                        "code": 11000,
                        "errmsg": message,
                        "op": documents[index],
                    }
                ],
                "writeConcernErrors": [],
                "nInserted": index,
                "nUpserted": 0,
                "nMatched": 0,
                "nModified": 0,
                "nRemoved": 0,
                "upserted": [],
            }
            raise BulkWriteError(result)

        return InsertManyResult(result)

    def replace_one(
        self,
        filter,
        replacement,
        upsert=False,
        bypass_document_validation=False,
        *args,
        **kwargs
    ):
        """ """
        validate_is_mapping("filter", filter)
        validate_ok_for_replace(replacement)
        validate_boolean("upsert", upsert)

        if bypass_document_validation:
            pass

        raw_result = {"n": 0, "nModified": 0}
        # updator = Updator(replacement)
        try:
            fw = next(self._internal_scan_query(filter))
        except StopIteration:
            if upsert:
                if "_id" not in replacement:
                    replacement["_id"] = bson.ObjectId()
                raw_result["upserted"] = replacement["_id"]
                raw_result["n"] = 1
                self._storage.write_one(self, replacement, check_keys=False)
        else:
            raw_result["n"] = 1
            if fw.doc != replacement:
                replacement["_id"] = fw.doc["_id"]
                self._storage.update_one(self, replacement)
                raw_result["nModified"] = 1

        return UpdateResult(raw_result)

    def _internal_scan_query(self, query_spec):
        """An internal document generator for update"""
        queryfilter = QueryFilter(query_spec)
        documents = self._storage.query(MontyCursor(self), 0)
        first_matched = None
        for doc in documents:
            if queryfilter(doc):
                first_matched = queryfilter.fieldwalker
                break

        if first_matched:
            yield first_matched  # for try statement to test update or insert
            yield first_matched  # start update, yield again
            # continue iter documents(generator)
            for doc in documents:
                if queryfilter(doc):
                    yield queryfilter.fieldwalker

    def _internal_upsert(self, query_spec, updator, raw_result):
        """Internal document upsert"""
        doc_cls = self._database.codec_options.document_class

        def _remove_dollar_key(doc):
            if is_duckument_type(doc):
                new_doc = doc_cls()
                fields = list(doc.keys())
                for field in fields:
                    if field[:1] == "$" or "." in field:
                        continue
                    new_doc[field] = _remove_dollar_key(doc[field])
                return new_doc
            else:
                return doc

        document = _remove_dollar_key(deepcopy(query_spec))
        if "_id" not in document:
            document["_id"] = bson.ObjectId()
        raw_result["upserted"] = document["_id"]
        raw_result["n"] = 1

        fieldwalker = FieldWalker(document)
        updator(fieldwalker, do_insert=True)
        self._storage.write_one(self, fieldwalker.doc)

    def _no_id_update(self, updator, filter=None):
        id_operator = updator.operations.get("_id")
        doc_id = (filter or {}).get("_id")
        if id_operator and id_operator._keep() != doc_id:
            msg = (
                "Performing an update on the path '_id' would "
                "modify the immutable field '_id'"
            )
            raise WriteError(msg, code=66)

    def update_one(
        self,
        filter,
        update,
        upsert=False,
        bypass_document_validation=False,
        array_filters=None,
        *args,
        **kwargs
    ):
        """ """
        validate_is_mapping("filter", filter)
        validate_ok_for_update(update)
        validate_list_or_none("array_filters", array_filters)
        validate_boolean("upsert", upsert)

        if bypass_document_validation:
            pass

        raw_result = {"n": 0, "nModified": 0}
        updator = Updator(update, array_filters)
        self._no_id_update(updator, filter)
        try:
            fw = next(self._internal_scan_query(filter))
        except StopIteration:
            if upsert:
                self._internal_upsert(filter, updator, raw_result)
        else:
            self._no_id_update(updator)

            raw_result["n"] = 1
            if updator(fw):
                self._storage.update_one(self, fw.doc)
                raw_result["nModified"] = 1

        return UpdateResult(raw_result)

    def update_many(
        self,
        filter,
        update,
        upsert=False,
        bypass_document_validation=False,
        array_filters=None,
        *args,
        **kwargs
    ):
        """ """
        validate_is_mapping("filter", filter)
        validate_ok_for_update(update)
        validate_list_or_none("array_filters", array_filters)
        validate_boolean("upsert", upsert)

        if bypass_document_validation:
            pass

        raw_result = {"n": 0, "nModified": 0}
        updator = Updator(update, array_filters)
        scanner = self._internal_scan_query(filter)
        self._no_id_update(updator, filter)
        try:
            next(scanner)
        except StopIteration:
            if upsert:
                self._internal_upsert(filter, updator, raw_result)
        else:
            self._no_id_update(updator)

            @on_err_close(scanner)
            def update_docs():
                n, m = 0, 0
                for fieldwalker in scanner:
                    n += 1
                    if updator(fieldwalker):
                        m += 1
                        yield fieldwalker.doc
                raw_result["n"] = n
                raw_result["nModified"] = m

            self._storage.update_many(self, update_docs())

        return UpdateResult(raw_result)

    def delete_one(self, filter):
        raw_result = {"n": 0}

        queryfilter = QueryFilter(filter)
        storage = self._storage
        documents = storage.query(MontyCursor(self), 0)

        for doc in documents:
            if queryfilter(doc):
                storage.delete_one(self, doc["_id"])
                raw_result["n"] = 1
                break

        return DeleteResult(raw_result)

    def delete_many(self, filter):
        raw_result = {"n": 0}

        queryfilter = QueryFilter(filter)
        storage = self._storage
        documents = storage.query(MontyCursor(self), 0)

        doc_ids = set()
        for doc in documents:
            if queryfilter(doc):
                doc_ids.add(doc["_id"])
                raw_result["n"] += 1

        storage.delete_many(self, doc_ids)

        return DeleteResult(raw_result)

    def find(self, *args, **kwargs):
        # return a cursor
        return MontyCursor(self, *args, **kwargs)

    def find_one(self, filter=None, *args, **kwargs):
        """ """
        if filter is not None and not isinstance(filter, abc.Mapping):
            filter = {"_id": filter}

        cursor = self.find(filter, *args, **kwargs)
        for result in cursor.limit(-1):
            return result
        return None

    def count(self, filter=None, **kwargs):
        warnings.warn(
            "count is deprecated. Use Collection.count_documents instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.count_documents(filter, **kwargs)

    def count_documents(self, filter, **kwargs):
        cursor = MontyCursor(self, filter=filter, **kwargs)
        return len(list(cursor))

    def distinct(self, key, filter=None, **kwargs):
        """ """
        if not isinstance(key, string_types):
            raise TypeError(
                "key must be an instance of %s" % (string_types.__name__,)
            )

        result = list()

        def get_value(doc):
            fieldwalker = FieldWalker(doc)
            fieldvalues = fieldwalker.go(key).get().value
            res = list()
            for v in fieldvalues.iter_flat():
                weighted = Weighted(v)
                if weighted not in result:
                    res.append(weighted)
            return res

        documents = self._storage.query(MontyCursor(self), 0)

        if filter:
            queryfilter = QueryFilter(filter)
            for doc in documents:
                if queryfilter(doc):
                    result += get_value(doc)
        else:
            for doc in documents:
                result += get_value(doc)

        return [weighted.value for weighted in sorted(result)]

    def drop(self):
        self._database.drop_collection(self._name)

    def save(self, to_save, *args, **kwargs):
        # DEPRECATED
        if "_id" in to_save:
            self.replace_one(
                {"_id": to_save["_id"]}, to_save, upsert=True, *args, **kwargs
            )
        else:
            self.insert_one(to_save, *args, **kwargs)

    def create_index(self, *args, **kwargs):
        """Not functioning, currently only exists for mongoengine support.
        """