From 74563a4bbf76e96678af813a0a53e24bc65cd56f Mon Sep 17 00:00:00 2001 From: gitovska Date: Tue, 19 May 2026 19:31:54 +0200 Subject: [PATCH] SNOW-3538315: add type annotations to DataFrame telemetry decorators --- src/snowflake/snowpark/_internal/telemetry.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/snowflake/snowpark/_internal/telemetry.py b/src/snowflake/snowpark/_internal/telemetry.py index 80f16d395e..da646b7216 100644 --- a/src/snowflake/snowpark/_internal/telemetry.py +++ b/src/snowflake/snowpark/_internal/telemetry.py @@ -9,7 +9,7 @@ from enum import Enum, unique from logging import getLogger import time -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, TypeVar from snowflake.connector import SnowflakeConnection from snowflake.connector.telemetry import ( @@ -48,6 +48,8 @@ _logger = getLogger(__name__) +_FuncT = TypeVar("_FuncT", bound=Callable[..., Any]) + @unique class TelemetryField(Enum): @@ -369,7 +371,7 @@ def wrap(*args, **kwargs): return wrap -def df_api_usage(func): +def df_api_usage(func: _FuncT) -> _FuncT: @functools.wraps(func) def wrap(*args, **kwargs): with ResourceUsageCollector() as resource_usage_collector: @@ -404,7 +406,7 @@ def wrap(*args, **kwargs): return wrap -def df_to_relational_group_df_api_usage(func): +def df_to_relational_group_df_api_usage(func: _FuncT) -> _FuncT: @functools.wraps(func) def wrap(*args, **kwargs): with ResourceUsageCollector() as resource_usage_collector: @@ -419,7 +421,7 @@ def wrap(*args, **kwargs): # For relational-grouped dataframe -def relational_group_df_api_usage(func): +def relational_group_df_api_usage(func: _FuncT) -> _FuncT: @functools.wraps(func) def wrap(*args, **kwargs): with ResourceUsageCollector() as resource_usage_collector: