Add type annotations to cqlengine models and query modules

cassandra/cqlengine/models.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 import re
+from typing import TYPE_CHECKING, TypeVar
 from warnings import warn
 
 from cassandra.cqlengine import CQLEngineException, ValidationError
@@ -329,6 +332,9 @@ def __delete__(self, instance):
                 raise AttributeError('cannot delete {0} columns'.format(self.column.column_name))
 
 
+M = TypeVar('M', bound='BaseModel')
+
+
 class BaseModel(object):
     """
     The base model class, don't inherit from this, inherit from Model, defined below
@@ -657,7 +663,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 +678,7 @@ def create(cls, **kwargs):
         return cls.objects.create(**kwargs)
 
     @classmethod
-    def all(cls):
+    def all(cls: type[M]) -> query.ModelQuerySet:
         """
         Returns a queryset representing all stored objects
 
@@ -681,7 +687,7 @@ def all(cls):
         return cls.objects.all()
 
     @classmethod
-    def filter(cls, *args, **kwargs):
+    def filter(cls: type[M], *args, **kwargs) -> query.ModelQuerySet:
         """
         Returns a queryset based on filter parameters.
 
@@ -690,15 +696,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 +713,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 +749,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

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, TypeVar
 from warnings import warn
 
 from cassandra.query import SimpleStatement, BatchType as CBatchType, BatchStatement
@@ -336,9 +339,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return
 
 
+if TYPE_CHECKING:
+    from cassandra.cqlengine.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 +535,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 +608,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 +625,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 +749,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 +790,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 +834,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 +887,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 +919,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 +975,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) \
@@ -1011,9 +1018,9 @@ def __eq__(self, q):
         return False
 
     def __ne__(self, q):
-        return not (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
@@ -1156,7 +1163,7 @@ def values_list(self, *fields, **kwargs):
         clone._flat_values_list = flat
         return clone
 
-    def ttl(self, ttl):
+    def ttl(self, ttl: int):
         """
         Sets the ttl (in seconds) for modified data.
 
@@ -1166,7 +1173,7 @@ def ttl(self, ttl):
         clone._ttl = ttl
         return clone
 
-    def timestamp(self, timestamp):
+    def timestamp(self, timestamp: datetime):
         """
         Allows for custom timestamps to be saved with the record.
         """