ref: stricter typing for a few utils modules (#93184)

<!-- Describe your PR here. -->

Author: asottile-sentry

pyproject.toml

@@ -406,19 +406,25 @@ module = [
     "sentry.utils.env",
     "sentry.utils.event",
     "sentry.utils.event_tracker",
+    "sentry.utils.forms",
     "sentry.utils.function_cache",
     "sentry.utils.geo",
     "sentry.utils.imports",
     "sentry.utils.iterators",
     "sentry.utils.javascript",
+    "sentry.utils.kafka",
     "sentry.utils.kvstore.*",
     "sentry.utils.lazy_service_wrapper",
     "sentry.utils.locking.*",
+    "sentry.utils.math",
+    "sentry.utils.memory",
+    "sentry.utils.metrics",
     "sentry.utils.migrations",
     "sentry.utils.numbers",
     "sentry.utils.otp",
     "sentry.utils.performance_issues.detectors.*",
     "sentry.utils.performance_issues.performance_detection",
+    "sentry.utils.projectflags",
     "sentry.utils.pubsub",
     "sentry.utils.redis",
     "sentry.utils.redis_metrics",
@@ -429,6 +435,7 @@ module = [
     "sentry.utils.services",
     "sentry.utils.sms",
     "sentry.utils.snowflake",
+    "sentry.utils.snuba_rpc",
     "sentry.utils.urls",
     "sentry.utils.uwsgi",
     "sentry.utils.warnings",

src/sentry/utils/forms.py

@@ -1,16 +1,34 @@
+from __future__ import annotations
+
 from collections.abc import Sequence
+from typing import TYPE_CHECKING, NotRequired, TypedDict
 
 from django import forms
 from django.forms import ChoiceField
 
+if TYPE_CHECKING:
+    from django.utils.functional import _StrPromise
+
+
+class _FieldConfig(TypedDict):
+    name: str
+    label: str | _StrPromise
+    placeholder: str | None
+    help: str | _StrPromise
+    required: bool
+    default: object
+    type: NotRequired[str]
+    choices: NotRequired[object]
 
-def field_to_config(name, field):
-    config = {
+
+def field_to_config(name: str, field: forms.Field) -> _FieldConfig:
+    config: _FieldConfig = {
         "name": name,
         "label": field.label or name.replace("_", " ").title(),
         "placeholder": field.widget.attrs.get("placeholder"),
         "help": field.help_text,
         "required": field.required,
+        "default": field.initial,
     }
     if isinstance(field, forms.URLField):
         config["type"] = "url"
@@ -33,13 +51,8 @@ def field_to_config(name, field):
     return config
 
 
-def form_to_config(form):
-    config = []
-    for name, field in form.base_fields.items():
-        row = field_to_config(name, field)
-        row["default"] = field.initial
-        config.append(row)
-    return config
+def form_to_config(form: forms.Form) -> list[_FieldConfig]:
+    return [field_to_config(name, field) for name, field in form.base_fields.items()]
 
 
 def set_field_choices(field: forms.Field, choices: Sequence[tuple[object, object]]) -> None:

src/sentry/utils/kafka.py

@@ -2,6 +2,9 @@
 import signal
 import time
 from threading import Thread
+from typing import Any
+
+from arroyo.processing.processor import StreamProcessor
 
 logger = logging.getLogger(__name__)
 
@@ -25,19 +28,21 @@ def delay_kafka_rebalance(configured_delay: int) -> None:
     time.sleep(sleep_secs)
 
 
-def delay_shutdown(processor, quantized_rebalance_delay_secs) -> None:
+def delay_shutdown(processor: StreamProcessor[Any], quantized_rebalance_delay_secs: int) -> None:
     if quantized_rebalance_delay_secs:
         delay_kafka_rebalance(quantized_rebalance_delay_secs)
 
     processor.signal_shutdown()
 
 
-def run_processor_with_signals(processor, quantized_rebalance_delay_secs: int | None = None):
+def run_processor_with_signals(
+    processor: StreamProcessor[Any], quantized_rebalance_delay_secs: int | None = None
+) -> None:
     if quantized_rebalance_delay_secs:
         # delay startup for quantization
         delay_kafka_rebalance(quantized_rebalance_delay_secs)
 
-    def handler(signum, frame):
+    def handler(signum: object, frame: object) -> None:
         # delay shutdown for quantization
         t = Thread(target=delay_shutdown, args=(processor, quantized_rebalance_delay_secs))
         t.start()

src/sentry/utils/math.py

@@ -1,20 +1,10 @@
+from __future__ import annotations
+
 import math
 from abc import ABC, abstractmethod
 
 
-def mean(values):
-    return sum(values) / len(values)
-
-
-def median(values):
-    values = sorted(values)
-    size = len(values)
-    if size % 2 == 1:
-        return values[int((size - 1) / 2)]
-    return (values[int(size / 2 - 1)] + values[int(size / 2)]) / 2
-
-
-def nice_int(x):
+def nice_int(x: float) -> int:
     """
     Round away from zero to the nearest "nice" number.
     """

src/sentry/utils/memory.py

@@ -1,15 +1,19 @@
+from __future__ import annotations
+
 import resource
+from collections.abc import Generator
 from contextlib import contextmanager
+from typing import Any
 
 from sentry.utils import metrics
 
 
-def get_rss_usage():
+def get_rss_usage() -> int:
     return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
 
 
 @contextmanager
-def track_memory_usage(metric, **kwargs):
+def track_memory_usage(metric: str, **kwargs: Any) -> Generator[None]:
     before = get_rss_usage()
     try:
         yield

src/sentry/utils/metrics.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import functools
 import logging
 import time
@@ -6,7 +8,7 @@
 from queue import Queue
 from random import random
 from threading import Thread
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import sentry_sdk
 from django.conf import settings
@@ -15,6 +17,9 @@
 from sentry.metrics.base import MetricsBackend, MutableTags, Tags
 from sentry.metrics.middleware import MiddlewareWrapper, add_global_tags, global_tags
 
+if TYPE_CHECKING:
+    from sentry.models.organization import Organization
+
 metrics_skip_all_internal = settings.SENTRY_METRICS_SKIP_ALL_INTERNAL
 metrics_skip_internal_prefixes = tuple(settings.SENTRY_METRICS_SKIP_INTERNAL_PREFIXES)
 
@@ -264,11 +269,11 @@ def event(
 def ensure_crash_rate_in_bounds(
     data: Any,
     request: Request,
-    organization,
+    organization: Organization,
     CRASH_RATE_METRIC_KEY: str,
     lower_bound: float = 0.0,
     upper_bound: float = 1.0,
-):
+) -> None:
     """
     Ensures that crash rate metric will always have value in expected bounds, and
     that invalid value is never returned to the customer by replacing all the invalid values with

src/sentry/utils/projectflags.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+from typing import Any
+
 from django.db.models import F
 from django.dispatch import Signal
 
 from sentry.models.project import Project
 
 
-def set_project_flag_and_signal(project: Project, flag_name: str, signal: Signal, **kwargs) -> bool:
+def set_project_flag_and_signal(
+    project: Project, flag_name: str, signal: Signal, **kwargs: Any
+) -> bool:
     """
     Helper function to set a project flag and send a signal.
     Returns True if the flag was set, False if it was already set.

src/sentry/utils/snuba_rpc.py

@@ -59,7 +59,7 @@ class MultiRpcResponse:
     timeseries_response: list[TimeSeriesResponse]
 
 
-def log_snuba_info(content):
+def log_snuba_info(content: str) -> None:
     if SNUBA_INFO_FILE:
         with open(SNUBA_INFO_FILE, "a") as file:
             file.writelines(content)