Fix mypy typing

grafana · Jul 31, 2024 · 3e16f5f · 3e16f5f
1 parent a2de1ba
commit 3e16f5f
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/engine/apps/email/models.py b/engine/apps/email/models.py
@@ -1,7 +1,13 @@
+import typing
 import uuid
 
 from django.db import models
 
+if typing.TYPE_CHECKING:
+    from apps.alerts.models import Alert, AlertGroup
+    from apps.base.models import UserNotificationPolicy
+    from apps.user_management.models import User
+
 
 class EmailMessageQuerySet(models.QuerySet):
     def create(self, **kwargs):
@@ -14,6 +20,11 @@ def create(self, **kwargs):
 
 
 class EmailMessage(models.Model):
+    represents_alert: "Alert"
+    represents_alert_group: "AlertGroup"
+    notification_policy: typing.Optional["UserNotificationPolicy"]
+    receiver: typing.Optional["User"]
+
     objects: models.Manager["EmailMessage"] = EmailMessageQuerySet.as_manager()
 
     message_uuid = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)