refactor: slightly improve typing for cqlengine

cassandra/cqlengine/models.py

@@ -14,6 +14,7 @@
 
 import logging
 import re
+from typing import Any, Type, TypeVar
 from warnings import warn
 
 from cassandra.cqlengine import CQLEngineException, ValidationError
@@ -85,7 +86,7 @@ class QuerySetDescriptor(object):
     it's declared on everytime it's accessed
     """
 
-    def __get__(self, obj, model):
+    def __get__(self, obj: Any, model: "BaseModel"):
         """ :rtype: ModelQuerySet """
         if model.__abstract__:
             raise CQLEngineException('cannot execute queries against abstract models')
@@ -328,6 +329,7 @@ def __delete__(self, instance):
             else:
                 raise AttributeError('cannot delete {0} columns'.format(self.column.column_name))
 
+M = TypeVar('M', bound='BaseModel')
 
 class BaseModel(object):
     """
@@ -340,7 +342,7 @@ class DoesNotExist(_DoesNotExist):
     class MultipleObjectsReturned(_MultipleObjectsReturned):
         pass
 
-    objects = QuerySetDescriptor()
+    objects: query.ModelQuerySet = QuerySetDescriptor()
     ttl = TTLDescriptor()
     consistency = ConsistencyDescriptor()
     iff = ConditionalDescriptor()
@@ -421,6 +423,7 @@ def __str__(self):
         return '{0} <{1}>'.format(self.__class__.__name__,
                                 ', '.join('{0}={1}'.format(k, getattr(self, k)) for k in self._primary_keys.keys()))
 
+
     @classmethod
     def _routing_key_from_values(cls, pk_values, protocol_version):
         return cls._key_serializer(pk_values, protocol_version)
@@ -657,7 +660,7 @@ def _as_dict(self):
         return values
 
     @classmethod
-    def create(cls, **kwargs):
+    def create(cls: Type[M], **kwargs) -> M:
         """
         Create an instance of this model in the database.
 
@@ -672,7 +675,7 @@ def create(cls, **kwargs):
         return cls.objects.create(**kwargs)
 
     @classmethod
-    def all(cls):
+    def all(cls: Type[M]) -> list[M]:
         """
         Returns a queryset representing all stored objects
 
@@ -681,7 +684,7 @@ def all(cls):
         return cls.objects.all()
 
     @classmethod
-    def filter(cls, *args, **kwargs):
+    def filter(cls: Type[M], *args, **kwargs):
         """
         Returns a queryset based on filter parameters.
 
@@ -690,15 +693,15 @@ def filter(cls, *args, **kwargs):
         return cls.objects.filter(*args, **kwargs)
 
     @classmethod
-    def get(cls, *args, **kwargs):
+    def get(cls: Type[M], *args, **kwargs) -> M:
         """
         Returns a single object based on the passed filter constraints.
 
         This is a pass-through to the model objects().:method:`~cqlengine.queries.get`.
         """
         return cls.objects.get(*args, **kwargs)
 
-    def timeout(self, timeout):
+    def timeout(self: M, timeout: float | None) -> M:
         """
         Sets a timeout for use in :meth:`~.save`, :meth:`~.update`, and :meth:`~.delete`
         operations
@@ -707,7 +710,7 @@ def timeout(self, timeout):
         self._timeout = timeout
         return self
 
-    def save(self):
+    def save(self: M) -> M:
         """
         Saves an object to the database.
 
@@ -743,7 +746,7 @@ def save(self):
 
         return self
 
-    def update(self, **values):
+    def update(self: M, **values) -> M:
         """
         Performs an update on the model instance. You can pass in values to set on the model
         for updating, or you can call without values to execute an update against any modified
@@ -834,9 +837,13 @@ def _class_get_connection(cls):
     def _inst_get_connection(self):
         return self._connection or self.__connection__
 
+    def __getitem__(self, s: slice | int) -> M | list[M]:
+        return self.objects.__getitem__(s)
+
     _get_connection = hybrid_classmethod(_class_get_connection, _inst_get_connection)
 
 
+
 class ModelMetaClass(type):
 
     def __new__(cls, name, bases, attrs):

cassandra/cqlengine/query.py

@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 from datetime import datetime, timedelta
 from functools import partial
 import time
+from typing import TYPE_CHECKING, ClassVar, Type, TypeVar
 from warnings import warn
 
 from cassandra.query import SimpleStatement, BatchType as CBatchType, BatchStatement
@@ -335,10 +338,12 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         return
 
+if TYPE_CHECKING:
+    from .models import M
 
 class AbstractQuerySet(object):
 
-    def __init__(self, model):
+    def __init__(self, model: Type[M]):
         super(AbstractQuerySet, self).__init__()
         self.model = model
 
@@ -528,7 +533,7 @@ def __iter__(self):
 
             idx += 1
 
-    def __getitem__(self, s):
+    def __getitem__(self, s: slice | int) -> M | list[M]:
         self._execute_query()
 
         if isinstance(s, slice):
@@ -601,7 +606,7 @@ def batch(self, batch_obj):
         clone._batch = batch_obj
         return clone
 
-    def first(self):
+    def first(self) -> M | None:
         try:
             return next(iter(self))
         except StopIteration:
@@ -618,7 +623,7 @@ def all(self):
         """
         return copy.deepcopy(self)
 
-    def consistency(self, consistency):
+    def consistency(self, consistency: int):
         """
         Sets the consistency level for the operation. See :class:`.ConsistencyLevel`.
 
@@ -742,7 +747,7 @@ def filter(self, *args, **kwargs):
 
         return clone
 
-    def get(self, *args, **kwargs):
+    def get(self, *args, **kwargs) -> M:
         """
         Returns a single instance matching this query, optionally with additional filter kwargs.
 
@@ -783,7 +788,7 @@ def _get_ordering_condition(self, colname):
 
         return colname, order_type
 
-    def order_by(self, *colnames):
+    def order_by(self, *colnames: str):
         """
         Sets the column(s) to be used for ordering
 
@@ -827,7 +832,7 @@ class Comment(Model):
         clone._order.extend(conditions)
         return clone
 
-    def count(self):
+    def count(self) -> int:
         """
         Returns the number of rows matched by this query.
 
@@ -880,7 +885,7 @@ class Automobile(Model):
 
         return clone
 
-    def limit(self, v):
+    def limit(self, v: int):
         """
         Limits the number of results returned by Cassandra. Use *0* or *None* to disable.
 
@@ -912,7 +917,7 @@ def limit(self, v):
         clone._limit = v
         return clone
 
-    def fetch_size(self, v):
+    def fetch_size(self, v: int):
         """
         Sets the number of rows that are fetched at a time.
 
@@ -968,15 +973,15 @@ def _only_or_defer(self, action, fields):
 
         return clone
 
-    def only(self, fields):
+    def only(self, fields: list[str]):
         """ Load only these fields for the returned query """
         return self._only_or_defer('only', fields)
 
-    def defer(self, fields):
+    def defer(self, fields: list[str]):
         """ Don't load these fields for the returned query """
         return self._only_or_defer('defer', fields)
 
-    def create(self, **kwargs):
+    def create(self, **kwargs) -> M:
         return self.model(**kwargs) \
             .batch(self._batch) \
             .ttl(self._ttl) \
@@ -1013,7 +1018,7 @@ def __eq__(self, q):
     def __ne__(self, q):
         return not (self != q)
 
-    def timeout(self, timeout):
+    def timeout(self, timeout: float | None):
         """
         :param timeout: Timeout for the query (in seconds)
         :type timeout: float or None
@@ -1064,6 +1069,7 @@ def _get_result_constructor(self):
         """
         return ResultObject
 
+T = TypeVar('T', bound=ModelQuerySet)
 
 class ModelQuerySet(AbstractQuerySet):
     """
@@ -1156,7 +1162,7 @@ def values_list(self, *fields, **kwargs):
         clone._flat_values_list = flat
         return clone
 
-    def ttl(self, ttl):
+    def ttl(self: T, ttl: int) -> T:
         """
         Sets the ttl (in seconds) for modified data.
 
@@ -1166,15 +1172,15 @@ def ttl(self, ttl):
         clone._ttl = ttl
         return clone
 
-    def timestamp(self, timestamp):
+    def timestamp(self: T, timestamp: datetime) -> T:
         """
         Allows for custom timestamps to be saved with the record.
         """
         clone = copy.deepcopy(self)
         clone._timestamp = timestamp
         return clone
 
-    def if_not_exists(self):
+    def if_not_exists(self: T) -> T:
         """
         Check the existence of an object before insertion.
 
@@ -1186,7 +1192,7 @@ def if_not_exists(self):
         clone._if_not_exists = True
         return clone
 
-    def if_exists(self):
+    def if_exists(self: T) -> T:
         """
         Check the existence of an object before an update or delete.