done

Author: jaykv

mongospecs/__init__.py

@@ -3,11 +3,11 @@
 from pymongo.collation import Collation
 from pymongo.operations import IndexModel
 
-from .empty import Empty
-from .ops import All, And, ElemMatch, Exists, In, Nor, Not, NotIn, Or, Size, SortBy, Type
-from .pagination import Page, Paginator
-from .query import Q
-from .se import MongoDecoder, MongoEncoder
+from .helpers.empty import Empty
+from .helpers.ops import All, And, ElemMatch, Exists, In, Nor, Not, NotIn, Or, Size, SortBy, Type
+from .helpers.pagination import Page, Paginator
+from .helpers.query import Q
+from .helpers.se import MongoDecoder, MongoEncoder
 
 __all__ = [
     # Queries

mongospecs/attrs.py

@@ -8,7 +8,6 @@
 from pymongo import MongoClient
 
 from .base import SpecBase, SubSpecBase
-from .types import SpecDocumentType
 from .helpers.empty import Empty
 from .helpers.se import MongoEncoder, mongo_dec_hook
 

mongospecs/base.py

@@ -3,20 +3,19 @@
 import typing as t
 from copy import deepcopy
 
-from mongospecs.types import SpecBaseType, SubSpecBaseType
+from mongospecs.types import SubSpecBaseType
 
 from .mixins.crud import CrudMixin
 from .mixins.index import IndexManagementMixin
 from .mixins.integrity import IntegrityMixin
-from .mixins.query import QueryMixin
 from .mixins.session import SessionTransactionMixin
 from .mixins.signal import SignalMixin
 
 T = t.TypeVar("T")
 
 
-class SpecBase(CrudMixin, IndexManagementMixin, IntegrityMixin, SessionTransactionMixin, SignalMixin):
-    ...
+class SpecBase(CrudMixin, IndexManagementMixin, IntegrityMixin, SessionTransactionMixin, SignalMixin): ...
+
 
 class SubSpecBase(SubSpecBaseType):
     _parent: t.ClassVar[t.Any] = SpecBase
@@ -85,4 +84,3 @@ def _projection_to_paths(cls, root_key: str, projection: t.Mapping[str, t.Any])
             return True
 
         return sub_projection
-

mongospecs/helpers/__init__.py

@@ -31,7 +31,7 @@ def encode(obj: t.Any, enc_hook: t.Optional[t.Callable[[t.Any], t.Any]] = bson_e
     return bson.encode(msgspec.to_builtins(obj, enc_hook=enc_hook, builtin_types=(bson.ObjectId,)))
 
 
-def encode_spec(obj: SpecBase[t.Any]) -> bytes:
+def encode_spec(obj: SpecBase) -> bytes:
     return bson.encode(obj.to_json_type())
 
 

mongospecs/helpers/bson.py

@@ -2,9 +2,9 @@
 
 from pymongo import ASCENDING, DESCENDING
 
-from .base import to_refs
+from mongospecs.utils import deep_merge, to_refs
+
 from .query import Condition, Group, Q
-from .utils import deep_merge
 
 __all__ = [
     # Operators

mongospecs/helpers/ops.py

@@ -7,7 +7,9 @@
 from copy import deepcopy
 from dataclasses import dataclass
 
-from .base import SpecBase, to_refs
+from mongospecs.base import SpecBase
+from mongospecs.utils import to_refs
+
 from .query import Condition, Group
 
 __all__ = (
@@ -20,7 +22,7 @@
 
 
 T = t.TypeVar("T")
-TSpec = t.TypeVar("TSpec", bound=SpecBase[t.Any])
+TSpec = t.TypeVar("TSpec", bound=SpecBase)
 
 
 class InvalidPage(Exception):

mongospecs/helpers/pagination.py

@@ -5,14 +5,14 @@
 from copy import deepcopy
 
 from bson import BSON, ObjectId
-from pymongo import ASCENDING, MongoClient
+from pymongo import MongoClient
 from pymongo.collection import Collection
 from pymongo.database import Database
 from typing_extensions import Self
 
 from ..helpers.empty import Empty, EmptyObject
-from ..types import SpecDocumentType, Specs, RawDocuments, SpecsOrRawDocuments, SubSpecBaseType, SpecBaseType
-import typing as t
+from ..types import RawDocuments, SpecBaseType, SpecDocumentType, SpecsOrRawDocuments
+
 
 class MongoBaseMixin(SpecBaseType):
     _client: t.ClassVar[t.Optional[MongoClient[t.Any]]] = None
@@ -38,9 +38,8 @@ def __eq__(self, other: t.Any) -> bool:
 
         return self._id == other._id
 
-    def __lt__(self, other: t.Any) -> t.Any:
-        return self._id < other._id
-    
+    def __lt__(self, other: t.Any) -> bool:
+        return self._id < other._id  # type: ignore[no-any-return]
 
     @classmethod
     def get_collection(cls) -> Collection[t.Any]:
@@ -76,7 +75,7 @@ def with_options(cls, **options: t.Any) -> t.Generator[t.Any, t.Any, None]:
                 cls._collection_context = existing_context
 
     @classmethod
-    def _path_to_value(cls, path: str, parent_dict: t.MutableMapping[str, t.Any]) -> t.Any:
+    def _path_to_value(cls, path: str, parent_dict: SpecDocumentType) -> t.Any:
         """Return a value from a dictionary at the given path"""
         keys: list[str] = cls._path_to_keys(path)
 
@@ -97,7 +96,7 @@ def _path_to_keys(cls, path: str) -> list[str]:
         return path.split(".")
 
     @classmethod
-    def _ensure_specs(cls, documents: SpecsOrRawDocuments[SpecDocumentType]) -> Specs[SpecDocumentType]:
+    def _ensure_specs(cls, documents: SpecsOrRawDocuments) -> t.Sequence[t.Self]:
         """
         Ensure all items in a list are specs by converting those that aren't.
         """
@@ -110,7 +109,7 @@ def _ensure_specs(cls, documents: SpecsOrRawDocuments[SpecDocumentType]) -> Spec
         return specs
 
     @classmethod
-    def _apply_sub_specs(cls, documents: RawDocuments[SpecDocumentType], subs: dict[str, t.Any]) -> None:
+    def _apply_sub_specs(cls, documents: RawDocuments, subs: dict[str, t.Any]) -> None:
         """Convert embedded documents to sub-specs for one or more documents"""
 
         # Dereference each reference
@@ -242,7 +241,7 @@ def _flatten_projection(
         return flat_projection, references, subs
 
     @classmethod
-    def _dereference(cls, documents: RawDocuments[t.Any], references: dict[str, t.Any]) -> None:
+    def _dereference(cls, documents: RawDocuments, references: dict[str, t.Any]) -> None:
         """Dereference one or more documents"""
 
         # Dereference each reference
@@ -320,24 +319,3 @@ def _remove_keys(cls, parent_dict: dict[str, t.Any], paths: list[str]) -> None:
                 continue
 
             child_dict.pop(keys[-1], None)
-
-
-def to_refs(value: t.Any) -> t.Any:
-    """Convert all Spec instances within the given value to Ids"""
-    # Spec
-    if isinstance(value, SpecBaseType):
-        return getattr(value, "_id")
-
-    # SubSpec
-    elif isinstance(value, SubSpecBaseType):
-        return to_refs(value.to_dict())
-
-    # Lists
-    elif isinstance(value, (list, tuple)):
-        return [to_refs(v) for v in value]
-
-    # Dictionaries
-    elif isinstance(value, dict):
-        return {k: to_refs(v) for k, v in value.items()}
-
-    return value

mongospecs/mixins/base.py

@@ -1,7 +1,14 @@
-from mongospecs.mixins.query import QueryMixin
 import typing as t
+
 from blinker import signal
-from ..types import SpecDocumentType, FilterType
+from pymongo import UpdateOne
+from pymongo.collection import Collection
+
+from mongospecs.helpers.query import Condition, Group
+from mongospecs.mixins.query import QueryMixin
+from mongospecs.types import FilterType, SpecDocumentType, SpecsOrRawDocuments
+from mongospecs.utils import to_refs
+
 
 class CrudMixin(QueryMixin):
     # Operations
@@ -171,7 +178,7 @@ def reload(self, **kwargs: t.Any) -> None:
             setattr(self, field, spec[field])
 
     @classmethod
-    def insert_many(cls, documents: SpecsOrRawDocuments[SpecDocumentType], **kwargs: t.Any) -> Specs[t.Any]:
+    def insert_many(cls, documents: t.Sequence[SpecDocumentType], **kwargs: t.Any) -> t.Sequence[t.Self]:
         """Insert a list of documents"""
         # Ensure all documents have been converted to specs
         specs = cls._ensure_specs(documents)
@@ -201,7 +208,7 @@ def insert_many(cls, documents: SpecsOrRawDocuments[SpecDocumentType], **kwargs:
     @classmethod
     def update_many(
         cls,
-        documents: SpecsOrRawDocuments[SpecDocumentType],
+        documents: SpecsOrRawDocuments,
         *fields: t.Any,
         update_one_kwargs: t.Any = None,
         bulk_write_kwargs: t.Any = None,
@@ -249,9 +256,7 @@ def update_many(
         signal("updated").send(cls, specs=specs)
 
     @classmethod
-    def unset_many(
-        cls, documents: SpecsOrRawDocuments[SpecDocumentType], *fields: t.Any, **update_many_kwargs: t.Any
-    ) -> None:
+    def unset_many(cls, documents: SpecsOrRawDocuments, *fields: t.Any, **update_many_kwargs: t.Any) -> None:
         """Unset the given list of fields for given documents."""
 
         # Ensure all documents have been converted to specs
@@ -278,7 +283,7 @@ def unset_many(
         signal("updated").send(cls, specs=specs)
 
     @classmethod
-    def delete_many(cls, documents: SpecsOrRawDocuments[SpecDocumentType], **delete_many_kwargs: t.Any) -> None:
+    def delete_many(cls, documents: SpecsOrRawDocuments, **delete_many_kwargs: t.Any) -> None:
         """Delete multiple documents"""
 
         # Ensure all documents have been converted to specs
@@ -298,7 +303,7 @@ def delete_many(cls, documents: SpecsOrRawDocuments[SpecDocumentType], **delete_
 
         # Send deleted signal
         signal("deleted").send(cls, specs=specs)
-        
+
     def soft_delete(self, **update_one_kwargs: t.Any) -> None:
         """Soft delete this document by setting a deleted flag."""
         assert "_id" in self.to_dict(), "Can't delete documents without `_id`"

mongospecs/mixins/crud.py

@@ -1,22 +1,18 @@
-
-
 import typing as t
+
 from pymongo import ASCENDING
 
+from ..types import SpecDocumentType
 from .base import MongoBaseMixin
 
-class IndexManagementMixin(MongoBaseMixin):
 
+class IndexManagementMixin(MongoBaseMixin):
     @classmethod
-    def create_index(cls, keys: t.Union[str, list[str]], direction: t.Union[int, str] = ASCENDING, **kwargs: t.Any) -> str:
+    def create_index(cls, keys: t.Union[str, list[tuple[str, int]]], **kwargs: t.Any) -> str:
         """
         Create an index on the specified keys (a single key or a list of keys).
         """
-        if isinstance(keys, str):
-            index_keys = [(keys, direction)]
-        else:
-            index_keys = [(key, direction) for key in keys]
-
+        index_keys = [(keys, ASCENDING)] if isinstance(keys, str) else keys
         return cls.get_collection().create_index(index_keys, **kwargs)
 
     @classmethod
@@ -27,7 +23,7 @@ def drop_index(cls, index_name: str) -> None:
         cls.get_collection().drop_index(index_name)
 
     @classmethod
-    def list_indexes(cls) -> list[t.MutableMapping[str, t.Any]]:
+    def list_indexes(cls) -> list[SpecDocumentType]:
         """
         List all indexes on the collection.
         """