type fixes

posthog/hogql_queries/insights/funnels/base.py

@@ -1,6 +1,6 @@
 from abc import ABC
 from functools import cached_property
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 import uuid
 from posthog.clickhouse.materialized_columns.column import ColumnName
 from posthog.constants import BREAKDOWN_VALUES_LIMIT
@@ -26,6 +26,7 @@
     BreakdownType,
     EventsNode,
     FunnelExclusionActionsNode,
+    FunnelTimeToConvertResults,
     StepOrderValue,
 )
 from posthog.types import EntityNode, ExclusionEntityNode
@@ -265,7 +266,9 @@ def _get_breakdown_expr(self) -> ast.Expr:
         else:
             raise ValidationError(detail=f"Unsupported breakdown type: {breakdownType}")
 
-    def _format_results(self, results) -> List[Dict[str, Any]] | List[List[Dict[str, Any]]]:
+    def _format_results(
+        self, results
+    ) -> Union[FunnelTimeToConvertResults, List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
         breakdown = self.context.breakdown
 
         if not results or len(results) == 0:

posthog/hogql_queries/insights/funnels/funnel_time_to_convert.py

@@ -6,6 +6,7 @@
 from posthog.hogql_queries.insights.funnels.base import FunnelBase
 from posthog.hogql_queries.insights.funnels.funnel_query_context import FunnelQueryContext
 from posthog.hogql_queries.insights.funnels.utils import get_funnel_order_class
+from posthog.schema import FunnelTimeToConvertResults
 
 
 class FunnelTimeToConvert(FunnelBase):
@@ -17,11 +18,11 @@ def __init__(
 
         self.funnel_order = get_funnel_order_class(self.context.funnelsFilter)(context=self.context)
 
-    def _format_results(self, results: list) -> dict:
-        return {
-            "bins": [(bin_from_seconds, person_count) for bin_from_seconds, person_count, _ in results],
-            "average_conversion_time": results[0][2],
-        }
+    def _format_results(self, results: list) -> FunnelTimeToConvertResults:
+        return FunnelTimeToConvertResults(
+            bins=[[bin_from_seconds, person_count] for bin_from_seconds, person_count, _ in results],
+            average_conversion_time=results[0][2],
+        )
 
     def get_query(self) -> ast.SelectQuery:
         query, funnelsFilter = self.context.query, self.context.funnelsFilter