ref: add typing-only hints to dict fields for django-stubs

asottile-sentry · asottile-sentry · commit 5d9446b8d072 · 2024-07-22T12:07:49.000-04:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -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",
@@ -530,7 +529,6 @@ module = [
     "sentry.issues.constants",
     "sentry.issues.endpoints",
     "sentry.issues.endpoints.group_events",
-    "sentry.issues.endpoints.organization_activity",
     "sentry.issues.endpoints.organization_group_search_views",
     "sentry.issues.endpoints.organization_release_previous_commits",
     "sentry.issues.endpoints.organization_searches",
@@ -626,7 +624,6 @@ module = [
     "tests.sentry.issues",
     "tests.sentry.issues.endpoints",
     "tests.sentry.issues.endpoints.test_actionable_items",
-    "tests.sentry.issues.endpoints.test_organization_activity",
     "tests.sentry.issues.endpoints.test_organization_group_search_views",
     "tests.sentry.issues.endpoints.test_organization_searches",
     "tests.sentry.issues.endpoints.test_project_stacktrace_link",
diff --git a/src/sentry/api/endpoints/project_rules.py b/src/sentry/api/endpoints/project_rules.py
@@ -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,
@@ -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
diff --git a/src/sentry/data_export/models.py b/src/sentry/data_export/models.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import Any
 
 import orjson
 from django.conf import settings
@@ -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:
diff --git a/src/sentry/discover/models.py b/src/sentry/discover/models.py
@@ -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
@@ -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)
diff --git a/src/sentry/models/activity.py b/src/sentry/models/activity.py
@@ -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()
 
diff --git a/src/sentry/models/authidentity.py b/src/sentry/models/authidentity.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Collection
 from typing import Any
 
@@ -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)
diff --git a/src/sentry/models/authprovider.py b/src/sentry/models/authprovider.py
@@ -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)
diff --git a/src/sentry/models/dashboard.py b/src/sentry/models/dashboard.py
@@ -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
 
diff --git a/src/sentry/models/debugfile.py b/src/sentry/models/debugfile.py
@@ -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()
diff --git a/src/sentry/models/files/abstractfile.py b/src/sentry/models/files/abstractfile.py
@@ -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
@@ -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
@@ -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)
 
@@ -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
 
diff --git a/src/sentry/models/group.py b/src/sentry/models/group.py
@@ -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)
diff --git a/src/sentry/models/grouplink.py b/src/sentry/models/grouplink.py
@@ -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
@@ -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()
diff --git a/src/sentry/models/groupsnooze.py b/src/sentry/models/groupsnooze.py
@@ -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
@@ -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",))
@@ -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
 
@@ -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
 
diff --git a/src/sentry/models/grouptombstone.py b/src/sentry/models/grouptombstone.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import Any
 
 from django.db import models
 
@@ -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:
diff --git a/src/sentry/models/identity.py b/src/sentry/models/identity.py
@@ -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)
 
@@ -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)
diff --git a/src/sentry/models/integrations/integration.py b/src/sentry/models/integrations/integration.py
@@ -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
     )
diff --git a/src/sentry/models/integrations/sentry_app_component.py b/src/sentry/models/integrations/sentry_app_component.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import MutableMapping
 from typing import Any
 
@@ -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"
diff --git a/src/sentry/models/notificationmessage.py b/src/sentry/models/notificationmessage.py
@@ -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
 
@@ -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.
diff --git a/src/sentry/models/organizationonboardingtask.py b/src/sentry/models/organizationonboardingtask.py
@@ -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
@@ -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]
 
diff --git a/src/sentry/models/projectownership.py b/src/sentry/models/projectownership.py
@@ -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)
diff --git a/src/sentry/models/repository.py b/src/sentry/models/repository.py
@@ -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
     )
diff --git a/src/sentry/models/rule.py b/src/sentry/models/rule.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Sequence
 from enum import Enum, IntEnum
 from typing import Any, ClassVar, Self
@@ -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")),
diff --git a/src/sentry/ownership/grammar.py b/src/sentry/ownership/grammar.py
diff --git a/src/sentry/testutils/factories.py b/src/sentry/testutils/factories.py
diff --git a/src/sentry/testutils/fixtures.py b/src/sentry/testutils/fixtures.py
