Skip to content

Commit

Permalink
ref: add typing-only hints to dict fields for django-stubs
Browse files Browse the repository at this point in the history
  • Loading branch information
asottile-sentry committed Jul 22, 2024
1 parent 7d304d1 commit 3bee4b4
Show file tree
Hide file tree
Showing 25 changed files with 80 additions and 55 deletions.
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,6 @@ module = [
"sentry.plugins.bases.notify",
"sentry.plugins.config",
"sentry.plugins.endpoints",
"sentry.plugins.providers.repository",
"sentry.receivers.releases",
"sentry.release_health.metrics_sessions_v2",
"sentry.replays.endpoints.project_replay_clicks_index",
Expand Down
14 changes: 6 additions & 8 deletions src/sentry/api/endpoints/project_rules.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,12 +81,12 @@ def __init__(
"""
rule.data will supersede rule_data if passed in
"""
self._project_id: int = project_id
self._rule_data: dict[Any, Any] = rule.data if rule else rule_data
self._rule_id: int | None = rule_id
self._rule: Rule | None = rule
self._project_id = project_id
self._rule_data = rule.data if rule else rule_data or {}
self._rule_id = rule_id
self._rule = rule

self._keys_to_check: set[str] = self._get_keys_to_check()
self._keys_to_check = self._get_keys_to_check()

self._matcher_funcs_by_key: dict[str, Callable[[Rule, str], MatcherResult]] = {
self.ENVIRONMENT_KEY: self._environment_matcher,
Expand All @@ -99,9 +99,7 @@ def _get_keys_to_check(self) -> set[str]:
Some keys are ignored as they are not part of the logic.
Some keys are required to check, and are added on top.
"""
keys_to_check: set[str] = {
key for key in list(self._rule_data.keys()) if key not in self.EXCLUDED_FIELDS
}
keys_to_check = {key for key in self._rule_data if key not in self.EXCLUDED_FIELDS}
keys_to_check.update(self.SPECIAL_FIELDS)

return keys_to_check
Expand Down
5 changes: 4 additions & 1 deletion src/sentry/data_export/models.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
from __future__ import annotations

import logging
from typing import Any

import orjson
from django.conf import settings
Expand Down Expand Up @@ -39,7 +42,7 @@ class ExportedData(Model):
date_finished = models.DateTimeField(null=True)
date_expired = models.DateTimeField(null=True, db_index=True)
query_type = BoundedPositiveIntegerField(choices=ExportQueryType.as_choices())
query_info = JSONField()
query_info: models.Field[dict[str, Any], dict[str, Any]] = JSONField()

@property
def status(self) -> ExportStatus:
Expand Down
6 changes: 4 additions & 2 deletions src/sentry/discover/models.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

from enum import Enum
from typing import ClassVar
from typing import Any, ClassVar

from django.db import models, router, transaction
from django.db.models import Q, UniqueConstraint
Expand Down Expand Up @@ -89,7 +91,7 @@ class DiscoverSavedQuery(Model):
organization = FlexibleForeignKey("sentry.Organization")
created_by_id = HybridCloudForeignKey("sentry.User", null=True, on_delete="SET_NULL")
name = models.CharField(max_length=255)
query = JSONField()
query: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
version = models.IntegerField(null=True)
date_created = models.DateTimeField(auto_now_add=True)
date_updated = models.DateTimeField(auto_now=True)
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/activity.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ class Activity(Model):
# if the user is not set, it's assumed to be the system
user_id = HybridCloudForeignKey(settings.AUTH_USER_MODEL, null=True, on_delete="SET_NULL")
datetime = models.DateTimeField(default=timezone.now)
data: models.Field[dict[str, Any], dict[str, Any]] = GzippedDictField(null=True)
data: models.Field[dict[str, Any] | None, dict[str, Any]] = GzippedDictField(null=True)

objects: ClassVar[ActivityManager] = ActivityManager()

Expand Down
4 changes: 3 additions & 1 deletion src/sentry/models/authidentity.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from collections.abc import Collection
from typing import Any

Expand Down Expand Up @@ -25,7 +27,7 @@ class AuthIdentity(ReplicatedControlModel):
user = FlexibleForeignKey(settings.AUTH_USER_MODEL)
auth_provider = FlexibleForeignKey("sentry.AuthProvider")
ident = models.CharField(max_length=128)
data = JSONField()
data: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
last_verified = models.DateTimeField(default=timezone.now)
last_synced = models.DateTimeField(default=timezone.now)
date_added = models.DateTimeField(default=timezone.now)
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/authprovider.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ class AuthProvider(ReplicatedControlModel):

organization_id = HybridCloudForeignKey("sentry.Organization", on_delete="cascade", unique=True)
provider = models.CharField(max_length=128)
config = JSONField()
config: models.Field[dict[str, Any], dict[str, Any]] = JSONField()

date_added = models.DateTimeField(default=timezone.now)
sync_time = BoundedPositiveIntegerField(null=True)
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/dashboard.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ class Dashboard(Model):
visits = BoundedBigIntegerField(null=True, default=1)
last_visited = models.DateTimeField(null=True, default=timezone.now)
projects = models.ManyToManyField("sentry.Project", through=DashboardProject)
filters = JSONField(null=True)
filters: models.Field[dict[str, Any] | None, dict[str, Any] | None] = JSONField(null=True)

MAX_WIDGETS = 30

Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/debugfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ class ProjectDebugFile(Model):
project_id = BoundedBigIntegerField(null=True)
debug_id = models.CharField(max_length=64, db_column="uuid")
code_id = models.CharField(max_length=64, null=True)
data = JSONField(null=True)
data: models.Field[dict[str, Any] | None, dict[str, Any] | None] = JSONField(null=True)
date_accessed = models.DateTimeField(default=timezone.now)

objects: ClassVar[ProjectDebugFileManager] = ProjectDebugFileManager()
Expand Down
7 changes: 4 additions & 3 deletions src/sentry/models/files/abstractfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import tempfile
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha1
from typing import ClassVar
from typing import Any, ClassVar

import sentry_sdk
from django.core.files.base import ContentFile
Expand All @@ -20,6 +20,7 @@
from sentry.celery import SentryTask
from sentry.db.models import BoundedPositiveIntegerField, JSONField, Model
from sentry.models.files.abstractfileblob import AbstractFileBlob
from sentry.models.files.abstractfileblobindex import AbstractFileBlobIndex
from sentry.models.files.utils import DEFAULT_BLOB_SIZE, AssembleChecksumMismatch, nooplogger
from sentry.utils import metrics
from sentry.utils.db import atomic_transaction
Expand Down Expand Up @@ -202,7 +203,7 @@ class AbstractFile(Model):
name = models.TextField()
type = models.CharField(max_length=64)
timestamp = models.DateTimeField(default=timezone.now, db_index=True)
headers = JSONField()
headers: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
size = BoundedPositiveIntegerField(null=True)
checksum = models.CharField(max_length=40, null=True, db_index=True)

Expand All @@ -212,7 +213,7 @@ class Meta:
# abstract
# XXX: uses `builtins.type` to avoid clash with `type` local
FILE_BLOB_MODEL: ClassVar[builtins.type[AbstractFileBlob]]
FILE_BLOB_INDEX_MODEL: ClassVar[builtins.type[Model]]
FILE_BLOB_INDEX_MODEL: ClassVar[builtins.type[AbstractFileBlobIndex]]
DELETE_UNREFERENCED_BLOB_TASK: ClassVar[SentryTask]
blobs: models.ManyToManyField

Expand Down
4 changes: 3 additions & 1 deletion src/sentry/models/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -565,7 +565,9 @@ class Group(Model):
score = BoundedIntegerField(default=0)
# deprecated, do not use. GroupShare has superseded
is_public = models.BooleanField(default=False, null=True)
data: models.Field[dict[str, Any], dict[str, Any]] = GzippedDictField(blank=True, null=True)
data: models.Field[dict[str, Any] | None, dict[str, Any]] = GzippedDictField(
blank=True, null=True
)
short_id = BoundedBigIntegerField(null=True)
type = BoundedPositiveIntegerField(default=ErrorGroupType.type_id, db_index=True)
priority = models.PositiveSmallIntegerField(null=True)
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/models/grouplink.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING, ClassVar
from typing import TYPE_CHECKING, Any, ClassVar

from django.db import models
from django.db.models import QuerySet
Expand Down Expand Up @@ -71,7 +71,7 @@ class LinkedType:
default=Relationship.references,
choices=((Relationship.resolves, _("Resolves")), (Relationship.references, _("Linked"))),
)
data = JSONField()
data: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
datetime = models.DateTimeField(default=timezone.now, db_index=True)

objects: ClassVar[GroupLinkManager] = GroupLinkManager()
Expand Down
11 changes: 6 additions & 5 deletions src/sentry/models/groupsnooze.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from datetime import timedelta
from typing import TYPE_CHECKING, ClassVar, Self
from typing import TYPE_CHECKING, Any, ClassVar, Self

from django.db import models
from django.db.models.signals import post_delete, post_save
Expand Down Expand Up @@ -53,7 +53,7 @@ class GroupSnooze(Model):
window = BoundedPositiveIntegerField(null=True)
user_count = BoundedPositiveIntegerField(null=True)
user_window = BoundedPositiveIntegerField(null=True)
state = JSONField(null=True)
state: models.Field[dict[str, Any] | None, dict[str, Any] | None] = JSONField(null=True)
actor_id = BoundedPositiveIntegerField(null=True)

objects: ClassVar[BaseManager[Self]] = BaseManager(cache_fields=("group",))
Expand Down Expand Up @@ -87,6 +87,7 @@ def is_valid(
return False
else:
times_seen = group.times_seen_with_pending if use_pending_data else group.times_seen
assert self.state is not None
if self.count <= times_seen - self.state["times_seen"]:
return False

Expand Down Expand Up @@ -200,10 +201,10 @@ def test_user_rates(self) -> bool:

def test_user_counts(self, group: Group) -> bool:
cache_key = f"groupsnooze:v1:{self.id}:test_user_counts:events_seen_counter"
try:
users_seen = self.state["users_seen"]
except (KeyError, TypeError):
if self.state is None:
users_seen = 0
else:
users_seen = self.state.get("users_seen", 0)

threshold = self.user_count + users_seen

Expand Down
7 changes: 6 additions & 1 deletion src/sentry/models/grouptombstone.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
from __future__ import annotations

import logging
from typing import Any

from django.db import models

Expand Down Expand Up @@ -29,7 +32,9 @@ class GroupTombstone(Model):
)
message = models.TextField()
culprit = models.CharField(max_length=MAX_CULPRIT_LENGTH, blank=True, null=True)
data = GzippedDictField(blank=True, null=True)
data: models.Field[dict[str, Any] | None, dict[str, Any]] = GzippedDictField(
blank=True, null=True
)
actor_id = BoundedPositiveIntegerField(null=True)

class Meta:
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/models/identity.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ class IdentityProvider(Model):
__relocation_scope__ = RelocationScope.Excluded

type = models.CharField(max_length=64)
config = JSONField()
config: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
date_added = models.DateTimeField(default=timezone.now, null=True)
external_id = models.CharField(max_length=64, null=True)

Expand Down Expand Up @@ -197,7 +197,7 @@ class Identity(Model):
idp = FlexibleForeignKey("sentry.IdentityProvider")
user = FlexibleForeignKey(settings.AUTH_USER_MODEL)
external_id = models.TextField()
data = JSONField()
data: models.Field[dict[str, Any], dict[str, Any]] = JSONField()
status = BoundedPositiveIntegerField(default=IdentityStatus.UNKNOWN)
scopes = ArrayField()
date_verified = models.DateTimeField(default=timezone.now)
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/integrations/integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ class Integration(DefaultFieldsModel):
# metadata might be used to store things like credentials, but it should NOT
# be used to store organization-specific information, as an Integration
# instance can be shared by multiple organizations
metadata = JSONField(default=dict)
metadata: models.Field[dict[str, Any], dict[str, Any]] = JSONField(default=dict)
status = BoundedPositiveIntegerField(
default=ObjectStatus.ACTIVE, choices=ObjectStatus.as_choices(), null=True
)
Expand Down
4 changes: 3 additions & 1 deletion src/sentry/models/integrations/sentry_app_component.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from collections.abc import MutableMapping
from typing import Any

Expand All @@ -17,7 +19,7 @@ class SentryAppComponent(Model):
uuid = UUIDField(unique=True, auto_add=True)
sentry_app = FlexibleForeignKey("sentry.SentryApp", related_name="components")
type = models.CharField(max_length=64)
schema = JSONField()
schema: models.Field[dict[str, Any], dict[str, Any]] = JSONField()

class Meta:
app_label = "sentry"
Expand Down
8 changes: 6 additions & 2 deletions src/sentry/models/notificationmessage.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
from django.db.models import DateTimeField, IntegerField, Q
from __future__ import annotations

from typing import Any

from django.db.models import DateTimeField, Field, IntegerField, Q
from django.db.models.constraints import CheckConstraint, UniqueConstraint
from django.utils import timezone

Expand Down Expand Up @@ -36,7 +40,7 @@ class NotificationMessage(Model):

# Related information regarding failed notifications.
# Leveraged to help give the user visibility into notifications that are consistently failing.
error_details = JSONField(null=True)
error_details: Field[dict[str, Any] | None, dict[str, Any] | None] = JSONField(null=True)
error_code = IntegerField(null=True, db_index=True)

# Resulting identifier from the vendor that can be leveraged for future interaction with the notification.
Expand Down
7 changes: 4 additions & 3 deletions src/sentry/models/organizationonboardingtask.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import ClassVar
from typing import Any, ClassVar

from django.conf import settings
from django.core.cache import cache
Expand Down Expand Up @@ -101,9 +101,10 @@ class AbstractOnboardingTask(Model):
completion_seen = models.DateTimeField(null=True)
date_completed = models.DateTimeField(default=timezone.now)
project = FlexibleForeignKey("sentry.Project", db_constraint=False, null=True)
data = JSONField() # INVITE_MEMBER { invited_member: user.id }
# INVITE_MEMBER { invited_member: user.id }
data: models.Field[dict[str, Any], dict[str, Any]] = JSONField()

# fields for typing
# abstract
TASK_LOOKUP_BY_KEY: dict[str, int]
SKIPPABLE_TASKS: frozenset[int]

Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/projectownership.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ class ProjectOwnership(Model):

project = FlexibleForeignKey("sentry.Project", unique=True)
raw = models.TextField(null=True)
schema = JSONField(null=True)
schema: models.Field[dict[str, Any] | None, dict[str, Any] | None] = JSONField(null=True)
fallthrough = models.BooleanField(default=True)
# Auto Assignment through Ownership Rules & Code Owners
auto_assignment = models.BooleanField(default=True)
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/repository.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ class Repository(Model, PendingDeletionMixin):
provider = models.CharField(max_length=64, null=True)
# The external_id is the id of the repo in the provider's system. (e.g. GitHub's repo id)
external_id = models.CharField(max_length=64, null=True)
config = JSONField(default=dict)
config: models.Field[dict[str, Any], dict[str, Any]] = JSONField(default=dict)
status = BoundedPositiveIntegerField(
default=ObjectStatus.ACTIVE, choices=ObjectStatus.as_choices(), db_index=True
)
Expand Down
4 changes: 3 additions & 1 deletion src/sentry/models/rule.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from collections.abc import Sequence
from enum import Enum, IntEnum
from typing import Any, ClassVar, Self
Expand Down Expand Up @@ -46,7 +48,7 @@ class Rule(Model):
environment_id = BoundedPositiveIntegerField(null=True)
label = models.CharField(max_length=256)
# `data` contain all the specifics of the rule - conditions, actions, frequency, etc.
data = GzippedDictField()
data: models.Field[dict[str, Any], dict[str, Any]] = GzippedDictField()
status = BoundedPositiveIntegerField(
default=ObjectStatus.ACTIVE,
choices=((ObjectStatus.ACTIVE, "Active"), (ObjectStatus.DISABLED, "Disabled")),
Expand Down
Loading

0 comments on commit 3bee4b4

Please sign in to comment.