refactor: Support type checking from cache_for decorator

mypy-baseline.txt

@@ -333,12 +333,14 @@ ee/billing/billing_manager.py:0: error: Incompatible types in assignment (expres
 posthog/models/property/util.py:0: error: Incompatible type for lookup 'pk': (got "str | int | list[str]", expected "str | int")  [misc]
 posthog/models/property/util.py:0: error: Argument 3 to "format_filter_query" has incompatible type "HogQLContext | None"; expected "HogQLContext"  [arg-type]
 posthog/models/property/util.py:0: error: Argument 3 to "format_cohort_subquery" has incompatible type "HogQLContext | None"; expected "HogQLContext"  [arg-type]
+posthog/models/property/util.py:0: error: Invalid index type "tuple[str, str]" for "dict[tuple[str, Literal['properties', 'group_properties', 'person_properties']], str]"; expected type "tuple[str, Literal['properties', 'group_properties', 'person_properties']]"  [index]
 posthog/models/property/util.py:0: error: Argument 1 to "append" of "list" has incompatible type "str | int"; expected "str"  [arg-type]
 posthog/models/property/util.py:0: error: Argument 1 to "append" of "list" has incompatible type "str | int"; expected "str"  [arg-type]
 posthog/models/property/util.py:0: error: Argument 1 to "append" of "list" has incompatible type "str | int"; expected "str"  [arg-type]
 posthog/api/utils.py:0: error: Incompatible types in assignment (expression has type "type[EventDefinition]", variable has type "type[EnterpriseEventDefinition]")  [assignment]
 posthog/api/utils.py:0: error: Argument 1 to "UUID" has incompatible type "int | str"; expected "str | None"  [arg-type]
 posthog/queries/trends/util.py:0: error: Argument 1 to "translate_hogql" has incompatible type "str | None"; expected "str"  [arg-type]
+posthog/queries/column_optimizer/foss_column_optimizer.py:0: error: Argument 1 to "get" of "dict" has incompatible type "tuple[str, str]"; expected "tuple[str, Literal['properties', 'group_properties', 'person_properties']]"  [arg-type]
 posthog/hogql/property.py:0: error: Incompatible type for lookup 'id': (got "str | int | list[str]", expected "str | int")  [misc]
 posthog/hogql/property.py:0: error: Incompatible type for lookup 'pk': (got "str | float", expected "str | int")  [misc]
 posthog/api/capture.py:0: error: Module has no attribute "utc"  [attr-defined]
@@ -599,6 +601,9 @@ posthog/api/organization_feature_flag.py:0: error: Invalid index type "str | Non
 posthog/api/organization_feature_flag.py:0: error: Invalid index type "str | None" for "dict[str, int]"; expected type "str"  [index]
 posthog/api/organization_feature_flag.py:0: error: Invalid index type "str | None" for "dict[str, int]"; expected type "str"  [index]
 posthog/api/notebook.py:0: error: Incompatible types in assignment (expression has type "int", variable has type "str | None")  [assignment]
+posthog/year_in_posthog/year_in_posthog.py:0: error: Item "None" of "dict[Any, Any] | None" has no attribute "get"  [union-attr]
+posthog/year_in_posthog/year_in_posthog.py:0: error: Argument 1 to "stats_for_user" has incompatible type "dict[Any, Any] | None"; expected "dict[Any, Any]"  [arg-type]
+posthog/year_in_posthog/year_in_posthog.py:0: error: Item "None" of "dict[Any, Any] | None" has no attribute "get"  [union-attr]
 posthog/warehouse/external_data_source/source.py:0: error: Incompatible types in assignment (expression has type "int", target has type "str")  [assignment]
 posthog/warehouse/external_data_source/source.py:0: error: Incompatible types in assignment (expression has type "int", target has type "str")  [assignment]
 posthog/warehouse/external_data_source/source.py:0: error: Incompatible types in assignment (expression has type "dict[str, Collection[str]]", variable has type "StripeSourcePayload")  [assignment]

posthog/cache_utils.py

@@ -1,7 +1,8 @@
+from dataclasses import dataclass, field
 import threading
-from datetime import timedelta
-from functools import wraps
-from typing import no_type_check, Any
+from collections.abc import Callable
+from datetime import datetime, timedelta
+from typing import Any, Generic, ParamSpec, TypeVar
 
 import orjson
 from rest_framework.utils.encoders import JSONEncoder
@@ -10,43 +11,54 @@
 
 from posthog.settings import TEST
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
+CacheKey = tuple[tuple[Any, ...], frozenset[tuple[Any, Any]]]
+
+
+@dataclass(slots=True)
+class CachedFunction(Generic[P, R]):
+    _fn: Callable[P, R]
+    _cache_time: timedelta
+    _background_refresh: bool = False
+
+    _cache: dict[CacheKey, tuple[datetime, R]] = field(default_factory=dict, init=False, repr=False)
+    _refreshing: dict[CacheKey, datetime | None] = field(default_factory=dict, init=False, repr=False)
+
+    def __call__(self, *args: P.args, use_cache: bool = not TEST, **kwargs: P.kwargs) -> R:
+        if not use_cache:
+            return self._fn(*args, **kwargs)
+
+        current_time = now()
+        key: CacheKey = (args, frozenset(sorted(kwargs.items())))
+
+        def refresh():
+            try:
+                value = self._fn(*args, **kwargs)
+                self._cache[key] = (now(), value)
+                self._refreshing[key] = None
+            except Exception:
+                self._refreshing[key] = None
+                raise
+
+        if key not in self._cache:
+            refresh()
+        elif current_time - self._cache[key][0] > self._cache_time:
+            if self._background_refresh:
+                if not self._refreshing.get(key):
+                    self._refreshing[key] = current_time
+                    t = threading.Thread(target=refresh)
+                    t.start()
+            else:
+                refresh()
 
-def cache_for(cache_time: timedelta, background_refresh=False):
-    def wrapper(fn):
-        @wraps(fn)
-        @no_type_check
-        def memoized_fn(*args, use_cache=not TEST, **kwargs):
-            if not use_cache:
-                return fn(*args, **kwargs)
-
-            current_time = now()
-            key = (args, frozenset(sorted(kwargs.items())))
+        return self._cache[key][1]
 
-            def refresh():
-                try:
-                    value = fn(*args, **kwargs)
-                    memoized_fn._cache[key] = (now(), value)
-                    memoized_fn._refreshing[key] = None
-                except Exception:
-                    memoized_fn._refreshing[key] = None
-                    raise
 
-            if key not in memoized_fn._cache:
-                refresh()
-            elif current_time - memoized_fn._cache[key][0] > cache_time:
-                if background_refresh:
-                    if not memoized_fn._refreshing.get(key):
-                        memoized_fn._refreshing[key] = current_time
-                        t = threading.Thread(target=refresh)
-                        t.start()
-                else:
-                    refresh()
-
-            return memoized_fn._cache[key][1]
-
-        memoized_fn._cache = {}
-        memoized_fn._refreshing = {}
-        return memoized_fn
+def cache_for(cache_time: timedelta, background_refresh=False) -> Callable[[Callable[P, R]], CachedFunction[P, R]]:
+    def wrapper(fn: Callable[P, R]) -> CachedFunction[P, R]:
+        return CachedFunction(fn, cache_time, background_refresh)
 
     return wrapper
 

posthog/clickhouse/migrations/0026_fix_materialized_window_and_session_ids.py

@@ -69,8 +69,8 @@ def materialize_session_and_window_id(database):
         # materialized the column or renamed the column, and then ran the 0004_...  async migration
         # before this migration runs.
         possible_old_column_names = {"mat_" + property_name}
-        current_materialized_column_name = materialized_columns.get(property_name, None)
-        if current_materialized_column_name != property_name:
+        current_materialized_column_name = materialized_columns.get((property_name, "properties"), None)
+        if current_materialized_column_name is not None and current_materialized_column_name != property_name:
             possible_old_column_names.add(current_materialized_column_name)
 
         for possible_old_column_name in possible_old_column_names: