Add mechanism for throttled warnings and use that for the ExperimentalWarning

Author: gabrys

src/neptune_query/internal/experimental.py

@@ -14,18 +14,15 @@
 # limitations under the License.
 
 import functools
-import warnings
 from typing import (
     Callable,
     ParamSpec,
     TypeVar,
 )
 
+from neptune_query.internal.warnings import throttled_warn
 from neptune_query.warnings import ExperimentalWarning
 
-# registry of functions already warned
-_warned_experimentals = set()
-
 T = ParamSpec("T")
 R = TypeVar("R")
 
@@ -38,14 +35,13 @@ def experimental(func: Callable[T, R]) -> Callable[T, R]:
 
     @functools.wraps(func)
     def wrapper(*args: T.args, **kwargs: T.kwargs) -> R:
-        if func not in _warned_experimentals:
-            warnings.warn(
+        throttled_warn(
+            ExperimentalWarning(
                 f"{func.__qualname__} is experimental and might change or be removed "
-                "in a future minor release. Use with caution in production code.",
-                category=ExperimentalWarning,
-                stacklevel=2,
-            )
-            _warned_experimentals.add(func)
+                "in a future minor release. Use with caution in production code."
+            ),
+            stacklevel=3,
+        )
         return func(*args, **kwargs)
 
     return wrapper

src/neptune_query/internal/warnings.py

@@ -0,0 +1,34 @@
+#
+# Copyright (c) 2025, Neptune Labs Sp. z o.o.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Type
+
+from neptune_query.warnings import ExperimentalWarning
+
+# registry of warnings that were already emitted with the (type, message) tuple
+WARNING_CLASSES_EMITTED_ONCE_PER_MESSAGE = (ExperimentalWarning,)
+
+_silence_warnings_msg: set[tuple[Type[Warning], str]] = set()
+
+
+def throttled_warn(warning: Warning, stacklevel: int = 2) -> None:
+    if isinstance(warning, WARNING_CLASSES_EMITTED_ONCE_PER_MESSAGE):
+        key = (type(warning), str(warning))
+        if key in _silence_warnings_msg:
+            return
+        _silence_warnings_msg.add(key)
+
+    warnings.warn(warning, stacklevel=stacklevel)