Browse Source

ref: fix typing for sentry.tsdb.dummy (#75475)

<!-- Describe your PR here. -->
anthony sottile 7 months ago
parent
commit
f528ab8e70
4 changed files with 56 additions and 48 deletions
  1. 0 1
      pyproject.toml
  2. 1 1
      src/sentry/tsdb/base.py
  3. 54 45
      src/sentry/tsdb/dummy.py
  4. 1 1
      src/sentry/tsdb/redis.py

+ 0 - 1
pyproject.toml

@@ -403,7 +403,6 @@ module = [
     "sentry.testutils.helpers.features",
     "sentry.testutils.helpers.notifications",
     "sentry.testutils.helpers.slack",
-    "sentry.tsdb.dummy",
     "sentry.tsdb.inmemory",
     "sentry.utils.auth",
     "sentry.utils.committers",

+ 1 - 1
src/sentry/tsdb/base.py

@@ -644,7 +644,7 @@ class BaseTSDB(Service):
         limit: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, str | int] | None = None,
-    ) -> dict[str, Iterable[dict[str, float]]]:
+    ) -> dict[str, Iterable[tuple[int, dict[str, float]]]]:
         """
         Retrieve the most frequently seen items in a frequency table for each
         interval in a series. (This is in contrast with ``get_most_frequent``,

+ 54 - 45
src/sentry/tsdb/dummy.py

@@ -1,9 +1,16 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Mapping, Sequence
 from datetime import datetime
+from typing import Any
 
 from sentry.tsdb.base import BaseTSDB, TSDBItem, TSDBKey, TSDBModel
 
 
+def _environment_ids(environment_ids: Iterable[int | None] | None) -> set[int | None]:
+    return (set(environment_ids) if environment_ids is not None else set()) | {None}
+
+
 class DummyTSDB(BaseTSDB):
     """
     A no-op time-series storage.
@@ -13,16 +20,10 @@ class DummyTSDB(BaseTSDB):
         self.validate_arguments([model], [environment_id])
 
     def merge(self, model, destination, sources, timestamp=None, environment_ids=None):
-        environment_ids = (set(environment_ids) if environment_ids is not None else set()).union(
-            [None]
-        )
-        self.validate_arguments([model], environment_ids)
+        self.validate_arguments([model], _environment_ids(environment_ids))
 
     def delete(self, models, keys, start=None, end=None, timestamp=None, environment_ids=None):
-        environment_ids = (set(environment_ids) if environment_ids is not None else set()).union(
-            [None]
-        )
-        self.validate_arguments(models, environment_ids)
+        self.validate_arguments(models, _environment_ids(environment_ids))
 
     def get_range(
         self,
@@ -38,7 +39,7 @@ class DummyTSDB(BaseTSDB):
         tenant_ids: dict[str, str | int] | None = None,
         referrer_suffix: str | None = None,
     ) -> dict[TSDBKey, list[tuple[int, int]]]:
-        self.validate_arguments([model], environment_ids if environment_ids is not None else [None])
+        self.validate_arguments([model], _environment_ids(environment_ids))
         _, series = self.get_optimal_rollup_series(start, end, rollup)
         return {k: [(ts, 0) for ts in series] for k in keys}
 
@@ -46,8 +47,15 @@ class DummyTSDB(BaseTSDB):
         self.validate_arguments([model], [environment_id])
 
     def get_distinct_counts_series(
-        self, model, keys: Sequence[int], start, end=None, rollup=None, environment_id=None
-    ):
+        self,
+        model: TSDBModel,
+        keys: Sequence[int],
+        start: datetime,
+        end: datetime | None = None,
+        rollup: int | None = None,
+        environment_id: int | None = None,
+        tenant_ids: dict[str, str | int] | None = None,
+    ) -> dict[int, list[tuple[int, Any]]]:
         self.validate_arguments([model], [environment_id])
         _, series = self.get_optimal_rollup_series(start, end, rollup)
         return {k: [(ts, 0) for ts in series] for k in keys}
@@ -69,26 +77,27 @@ class DummyTSDB(BaseTSDB):
         return {k: 0 for k in keys}
 
     def get_distinct_counts_union(
-        self, model, keys, start, end=None, rollup=None, environment_id=None
-    ):
+        self,
+        model: TSDBModel,
+        keys: list[int] | None,
+        start: datetime,
+        end: datetime | None = None,
+        rollup: int | None = None,
+        environment_id: int | None = None,
+        tenant_ids: dict[str, str | int] | None = None,
+    ) -> int:
         self.validate_arguments([model], [environment_id])
         return 0
 
     def merge_distinct_counts(
         self, model, destination, sources, timestamp=None, environment_ids=None
     ):
-        environment_ids = (set(environment_ids) if environment_ids is not None else set()).union(
-            [None]
-        )
-        self.validate_arguments([model], environment_ids)
+        self.validate_arguments([model], _environment_ids(environment_ids))
 
     def delete_distinct_counts(
         self, models, keys, start=None, end=None, timestamp=None, environment_ids=None
     ):
-        environment_ids = (set(environment_ids) if environment_ids is not None else set()).union(
-            [None]
-        )
-        self.validate_arguments(models, environment_ids)
+        self.validate_arguments(models, _environment_ids(environment_ids))
 
     def record_frequency_multi(
         self,
@@ -100,20 +109,29 @@ class DummyTSDB(BaseTSDB):
 
     def get_most_frequent(
         self,
-        model,
+        model: TSDBModel,
         keys: Sequence[TSDBKey],
-        start,
-        end=None,
-        rollup=None,
-        limit=None,
-        environment_id=None,
-    ):
+        start: datetime,
+        end: datetime | None = None,
+        rollup: int | None = None,
+        limit: int | None = None,
+        environment_id: int | None = None,
+        tenant_ids: dict[str, str | int] | None = None,
+    ) -> dict[TSDBKey, list[tuple[str, float]]]:
         self.validate_arguments([model], [environment_id])
         return {key: [] for key in keys}
 
     def get_most_frequent_series(
-        self, model, keys, start, end=None, rollup=None, limit=None, environment_id=None
-    ):
+        self,
+        model: TSDBModel,
+        keys: Iterable[str],
+        start: datetime,
+        end: datetime | None = None,
+        rollup: int | None = None,
+        limit: int | None = None,
+        environment_id: int | None = None,
+        tenant_ids: dict[str, str | int] | None = None,
+    ) -> dict[str, Iterable[tuple[int, dict[str, float]]]]:
         self.validate_arguments([model], [environment_id])
         rollup, series = self.get_optimal_rollup_series(start, end, rollup)
         return {key: [(timestamp, {}) for timestamp in series] for key in keys}
@@ -131,13 +149,10 @@ class DummyTSDB(BaseTSDB):
         self.validate_arguments([model], [environment_id])
         rollup, series = self.get_optimal_rollup_series(start, end, rollup)
 
-        results = {}
-        for key, members in items.items():
-            result = results[key] = []
-            for timestamp in series:
-                result.append((timestamp, {k: 0.0 for k in members}))
-
-        return results
+        return {
+            key: [(timestamp, {k: 0.0 for k in members}) for timestamp in series]
+            for key, members in items.items()
+        }
 
     def get_frequency_totals(
         self,
@@ -163,18 +178,12 @@ class DummyTSDB(BaseTSDB):
         timestamp: datetime | None = None,
         environment_ids: Iterable[int] | None = None,
     ) -> None:
-        environment_ids = list(
-            (set(environment_ids) if environment_ids is not None else set()).union([None])
-        )
-        self.validate_arguments([model], environment_ids)
+        self.validate_arguments([model], _environment_ids(environment_ids))
 
     def delete_frequencies(
         self, models, keys, start=None, end=None, timestamp=None, environment_ids=None
     ):
-        environment_ids = (set(environment_ids) if environment_ids is not None else set()).union(
-            [None]
-        )
-        self.validate_arguments(models, environment_ids)
+        self.validate_arguments(models, _environment_ids(environment_ids))
 
     def flush(self):
         pass

+ 1 - 1
src/sentry/tsdb/redis.py

@@ -878,7 +878,7 @@ class RedisTSDB(BaseTSDB):
         limit: int | None = None,
         environment_id: int | None = None,
         tenant_ids: dict[str, int | str] | None = None,
-    ) -> dict[str, Iterable[dict[str, float]]]:
+    ) -> dict[str, Iterable[tuple[int, dict[str, float]]]]:
         self.validate_arguments([model], [environment_id])
 
         if not self.enable_frequency_sketches: