diff --git a/posthog/hogql_queries/insights/funnels/base.py b/posthog/hogql_queries/insights/funnels/base.py index 0ec978f692b57..f742d17e767ff 100644 --- a/posthog/hogql_queries/insights/funnels/base.py +++ b/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: diff --git a/posthog/hogql_queries/insights/funnels/funnel_time_to_convert.py b/posthog/hogql_queries/insights/funnels/funnel_time_to_convert.py index afa38e1ee7382..38600dcc1fec6 100644 --- a/posthog/hogql_queries/insights/funnels/funnel_time_to_convert.py +++ b/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