Skip to content

Commit

Permalink
Improve type hints in hc.api.models
Browse files Browse the repository at this point in the history
  • Loading branch information
cuu508 committed Sep 8, 2023
1 parent a6ab647 commit 5746bbb
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 64 deletions.
98 changes: 58 additions & 40 deletions hc/api/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from datetime import datetime
from datetime import timedelta as td
from datetime import timezone
from typing import TypedDict
from typing import Any, TypedDict
from urllib.parse import urlencode
from zoneinfo import ZoneInfo

Expand All @@ -17,6 +17,7 @@
from django.core.mail import mail_admins
from django.core.signing import TimestampSigner
from django.db import models, transaction
from django.db.models import QuerySet
from django.http import HttpRequest
from django.urls import reverse
from django.utils.functional import cached_property
Expand All @@ -28,6 +29,7 @@
from hc.lib import emails
from hc.lib.date import month_boundaries
from hc.lib.s3 import get_object, put_object, remove_objects
from hc.lib.typealias import JSONDict

STATUSES = (("up", "Up"), ("down", "Down"), ("new", "New"), ("paused", "Paused"))
DEFAULT_TIMEOUT = td(days=1)
Expand All @@ -37,7 +39,7 @@
# max time between start and ping where we will consider both events related:
MAX_DURATION = td(hours=72)

TRANSPORTS = {
TRANSPORTS: dict[str, tuple[str, type[transports.Transport]]] = {
"apprise": ("Apprise", transports.Apprise),
"call": ("Phone Call", transports.Call),
"discord": ("Discord", transports.Discord),
Expand Down Expand Up @@ -70,6 +72,7 @@
"zulip": ("Zulip", transports.Zulip),
}


CHANNEL_KINDS = [(kind, label_cls[0]) for kind, label_cls in TRANSPORTS.items()]

PO_PRIORITIES = {
Expand All @@ -91,7 +94,7 @@
}


def isostring(dt) -> str | None:
def isostring(dt: datetime | None) -> str | None:
"""Convert the datetime to ISO 8601 format with no microseconds."""
return dt.replace(microsecond=0).isoformat() if dt else None

Expand Down Expand Up @@ -129,7 +132,7 @@ class CheckDict(TypedDict, total=False):


class DowntimeSummary(object):
def __init__(self, boundaries: list[datetime]):
def __init__(self, boundaries: list[datetime]) -> None:
self.boundaries = list(sorted(boundaries, reverse=True))
self.durations = [td() for _ in boundaries]
self.counts = [0 for _ in boundaries]
Expand Down Expand Up @@ -187,7 +190,7 @@ class Meta:
models.Index(fields=["project_id", "slug"], name="api_check_project_slug"),
]

def __str__(self):
def __str__(self) -> str:
return "%s (%d)" % (self.name or self.code, self.id)

def name_then_code(self) -> str:
Expand All @@ -214,7 +217,7 @@ def url(self) -> str | None:

return settings.PING_ENDPOINT + str(self.code)

def details_url(self, full=True) -> str:
def details_url(self, full: bool = True) -> str:
result = reverse("hc-details", args=[self.code])
return settings.SITE_ROOT + result if full else result

Expand Down Expand Up @@ -479,11 +482,13 @@ def prune(self) -> None:
pass

@property
def visible_pings(self):
def visible_pings(self) -> QuerySet["Ping"]:
threshold = self.n_pings - self.project.owner_profile.ping_log_limit
return self.ping_set.filter(n__gt=threshold)

def downtimes_by_boundary(self, boundaries: list[datetime]):
def downtimes_by_boundary(
self, boundaries: list[datetime]
) -> list[tuple[datetime, td | None, int | None]]:
"""Calculate downtime counts and durations for the given time intervals.
Returns a list of (datetime, downtime_in_secs, number_of_outages) tuples
Expand Down Expand Up @@ -528,7 +533,9 @@ def downtimes_by_boundary(self, boundaries: list[datetime]):
result.sort()
return result

def downtimes(self, months: int, tz: str):
def downtimes(
self, months: int, tz: str
) -> list[tuple[datetime, td | None, int | None]]:
boundaries = month_boundaries(months, tz)
return self.downtimes_by_boundary(boundaries)

Expand Down Expand Up @@ -615,7 +622,7 @@ def has_body(self) -> bool:
def get_body_bytes(self) -> bytes | None:
if self.body:
return self.body.encode()
if self.object_size:
if self.object_size and self.n:
return get_object(self.owner.code, self.n)
if self.body_raw:
return self.body_raw
Expand Down Expand Up @@ -674,9 +681,11 @@ def duration(self) -> td | None:


def json_property(kind: str, field: str) -> property:
def fget(instance):
def fget(instance: Channel) -> int | str:
assert instance.kind == kind
return instance.json[field]
v = instance.json[field]
assert isinstance(v, int) or isinstance(v, str)
return v

return property(fget)

Expand All @@ -702,7 +711,7 @@ class Channel(models.Model):
last_error = models.CharField(max_length=200, blank=True)
checks = models.ManyToManyField(Check)

def __str__(self):
def __str__(self) -> str:
if self.name:
return self.name
if self.kind == "email":
Expand Down Expand Up @@ -766,7 +775,7 @@ def send_signal_captcha_alert(self, challenge: str, raw: str) -> None:
"""
mail_admins(subject, message, html_message=html_message)

def send_signal_rate_limited_notice(self, message: str, plaintext: str):
def send_signal_rate_limited_notice(self, message: str, plaintext: str) -> None:
email = self.project.owner.email
ctx = {
"recipient": self.phone_number,
Expand All @@ -777,7 +786,7 @@ def send_signal_rate_limited_notice(self, message: str, plaintext: str):
emails.signal_rate_limited(email, ctx)

@property
def transport(self):
def transport(self) -> transports.Transport:
if self.kind not in TRANSPORTS:
raise NotImplementedError(f"Unknown channel kind: {self.kind}")

Expand Down Expand Up @@ -821,11 +830,11 @@ def icon_path(self) -> str:
return f"img/integrations/{self.kind}.png"

@property
def json(self):
def json(self) -> Any:
return json.loads(self.value)

@property
def po_priority(self):
def po_priority(self) -> str:
assert self.kind == "po"
parts = self.value.split("|")
prio = int(parts[1])
Expand Down Expand Up @@ -855,41 +864,49 @@ def up_webhook_spec(self) -> WebhookSpec:
cmd_up = json_property("shell", "cmd_up")

@property
def slack_team(self):
def slack_team(self) -> str | None:
assert self.kind == "slack"
if not self.value.startswith("{"):
return None

doc = json.loads(self.value)
if "team_name" in doc:
assert isinstance(doc["team_name"], str)
return doc["team_name"]

if "team" in doc:
assert isinstance(doc["team"]["name"], str)
return doc["team"]["name"]

return None

@property
def slack_channel(self):
def slack_channel(self) -> str | None:
assert self.kind == "slack"
if not self.value.startswith("{"):
return None

doc = json.loads(self.value)
return doc["incoming_webhook"]["channel"]
v = doc["incoming_webhook"]["channel"]
assert isinstance(v, str)
return v

@property
def slack_webhook_url(self):
def slack_webhook_url(self) -> str:
assert self.kind in ("slack", "mattermost")
if not self.value.startswith("{"):
return self.value

doc = json.loads(self.value)
return doc["incoming_webhook"]["url"]
v = doc["incoming_webhook"]["url"]
assert isinstance(v, str)
return v

@property
def discord_webhook_url(self):
def discord_webhook_url(self) -> str:
assert self.kind == "discord"
url = self.json["webhook"]["url"]

assert isinstance(url, str)
# Discord migrated to discord.com,
# and is dropping support for discordapp.com on 7 November 2020
if url.startswith("https://discordapp.com/"):
Expand Down Expand Up @@ -924,21 +941,22 @@ def update_telegram_id(self, new_chat_id) -> None:
self.save()

@property
def pd_service_key(self):
def pd_service_key(self) -> str:
assert self.kind == "pd"
if not self.value.startswith("{"):
return self.value

return self.json["service_key"]

@property
def pd_account(self):
def pd_account(self) -> str | None:
assert self.kind == "pd"
if self.value.startswith("{"):
return self.json.get("account")
return None

@property
def phone_number(self):
def phone_number(self) -> str:
assert self.kind in ("call", "sms", "whatsapp", "signal")
if self.value.startswith("{"):
return self.json["value"]
Expand All @@ -949,29 +967,29 @@ def phone_number(self):
trello_list_id = json_property("trello", "list_id")

@property
def trello_board_list(self):
def trello_board_list(self) -> tuple[str, str]:
assert self.kind == "trello"
doc = json.loads(self.value)
return doc["board_name"], doc["list_name"]

@property
def email_value(self):
def email_value(self) -> str:
assert self.kind == "email"
if not self.value.startswith("{"):
return self.value

return self.json["value"]

@property
def email_notify_up(self):
def email_notify_up(self) -> bool:
assert self.kind == "email"
if not self.value.startswith("{"):
return True

return self.json.get("up")

@property
def email_notify_down(self):
def email_notify_down(self) -> bool:
assert self.kind == "email"
if not self.value.startswith("{"):
return True
Expand All @@ -985,25 +1003,25 @@ def email_notify_down(self):
signal_notify_down = json_property("signal", "down")

@property
def sms_notify_up(self):
def sms_notify_up(self) -> bool:
assert self.kind == "sms"
return self.json.get("up", False)

@property
def sms_notify_down(self):
def sms_notify_down(self) -> bool:
assert self.kind == "sms"
return self.json.get("down", True)

@property
def opsgenie_key(self):
def opsgenie_key(self) -> str:
assert self.kind == "opsgenie"
if not self.value.startswith("{"):
return self.value

return self.json["key"]

@property
def opsgenie_region(self):
def opsgenie_region(self) -> str:
assert self.kind == "opsgenie"
if not self.value.startswith("{"):
return "us"
Expand All @@ -1016,7 +1034,7 @@ def opsgenie_region(self):
zulip_to = json_property("zulip", "to")

@property
def zulip_site(self):
def zulip_site(self) -> str:
assert self.kind == "zulip"
doc = json.loads(self.value)
if "site" in doc:
Expand All @@ -1028,12 +1046,12 @@ def zulip_site(self):
return "https://" + domain

@property
def zulip_topic(self):
def zulip_topic(self) -> str:
assert self.kind == "zulip"
return self.json.get("topic", "")

@property
def linenotify_token(self):
def linenotify_token(self) -> str:
assert self.kind == "linenotify"
return self.value

Expand All @@ -1046,12 +1064,12 @@ def linenotify_token(self):
ntfy_priority_up = json_property("ntfy", "priority_up")

@property
def ntfy_token(self):
def ntfy_token(self) -> str | None:
assert self.kind == "ntfy"
return self.json.get("token")

@property
def ntfy_priority_display(self):
def ntfy_priority_display(self) -> str:
return NTFY_PRIORITIES[self.ntfy_priority]


Expand Down
Loading

0 comments on commit 5746bbb

Please sign in to comment.