Skip to content

Commit 9a83826

Browse files
ref: improve typing of base_query_set (#72719)
<!-- Describe your PR here. -->
1 parent 01d69f1 commit 9a83826

File tree

2 files changed

+16
-11
lines changed

2 files changed

+16
-11
lines changed

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -521,6 +521,7 @@ module = [
521521
"sentry.api.helpers.source_map_helper",
522522
"sentry.buffer.*",
523523
"sentry.build.*",
524+
"sentry.db.models.manager.base_query_set",
524525
"sentry.eventstore.reprocessing.redis",
525526
"sentry.eventtypes.error",
526527
"sentry.grouping.component",

src/sentry/db/models/manager/base_query_set.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,22 @@
1-
import abc
2-
from typing import Any
1+
from __future__ import annotations
2+
3+
from typing import Any, Self
34

45
from django.core import exceptions
56
from django.core.exceptions import EmptyResultSet
67
from django.db import connections, router, transaction
78
from django.db.models import QuerySet, sql
89

10+
from sentry.db.models.manager import M
911
from sentry.signals import post_update
1012

1113

12-
class BaseQuerySet(QuerySet, abc.ABC):
13-
def __init__(self, *args, **kwargs):
14+
class BaseQuerySet(QuerySet[M]):
15+
def __init__(self, *args: Any, **kwargs: Any) -> None:
1416
super().__init__(*args, **kwargs)
1517
self._with_post_update_signal = False
1618

17-
def with_post_update_signal(self, enable: bool) -> "BaseQuerySet":
19+
def with_post_update_signal(self, enable: bool) -> Self:
1820
"""
1921
Enables sending a `post_update` signal after this queryset runs an update command. Note that this is less
2022
efficient than just running the update. To get the list of group ids affected, we first run the query to
@@ -24,12 +26,12 @@ def with_post_update_signal(self, enable: bool) -> "BaseQuerySet":
2426
qs._with_post_update_signal = enable
2527
return qs
2628

27-
def _clone(self) -> "BaseQuerySet":
29+
def _clone(self) -> Self:
2830
qs = super()._clone() # type: ignore[misc]
2931
qs._with_post_update_signal = self._with_post_update_signal
3032
return qs
3133

32-
def update_with_returning(self, returned_fields: list[str], **kwargs):
34+
def update_with_returning(self, returned_fields: list[str], **kwargs: Any) -> list[tuple[int]]:
3335
"""
3436
Copied and modified from `Queryset.update()` to support `RETURNING <returned_fields>`
3537
"""
@@ -77,7 +79,9 @@ def update_with_returning(self, returned_fields: list[str], **kwargs):
7779

7880
def update(self, **kwargs: Any) -> int:
7981
if self._with_post_update_signal:
80-
pk = self.model._meta.pk.name
82+
pk_field = self.model._meta.pk
83+
assert pk_field is not None
84+
pk = pk_field.name
8185
ids = [result[0] for result in self.update_with_returning([pk], **kwargs)]
8286
if ids:
8387
updated_fields = list(kwargs.keys())
@@ -86,17 +90,17 @@ def update(self, **kwargs: Any) -> int:
8690
else:
8791
return super().update(**kwargs)
8892

89-
def using_replica(self) -> "BaseQuerySet":
93+
def using_replica(self) -> Self:
9094
"""
9195
Use read replica for this query. Database router is expected to use the
9296
`replica=True` hint to make routing decision.
9397
"""
9498
return self.using(router.db_for_read(self.model, replica=True))
9599

96-
def defer(self, *args: Any, **kwargs: Any) -> "BaseQuerySet":
100+
def defer(self, *args: Any, **kwargs: Any) -> Self:
97101
raise NotImplementedError("Use ``values_list`` instead [performance].")
98102

99-
def only(self, *args: Any, **kwargs: Any) -> "BaseQuerySet":
103+
def only(self, *args: Any, **kwargs: Any) -> Self:
100104
# In rare cases Django can use this if a field is unexpectedly deferred. This
101105
# mostly can happen if a field is added to a model, and then an old pickle is
102106
# passed to a process running the new code. So if you see this error after a

0 commit comments

Comments
 (0)