diff --git a/cassandra/cqlengine/models.py b/cassandra/cqlengine/models.py index bc00001666..b3fce34d50 100644 --- a/cassandra/cqlengine/models.py +++ b/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,7 +696,7 @@ 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. @@ -698,7 +704,7 @@ def get(cls, *args, **kwargs): """ 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 diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py index afc7ceeef6..52ad090ac7 100644 --- a/cassandra/cqlengine/query.py +++ b/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. """