Custom warn once logic

src/neptune/new/internal/utils/deprecation.py

@@ -19,23 +19,40 @@
 
 from neptune.new.exceptions import NeptuneParametersCollision
 
-__all__ = ["deprecated", "deprecated_parameter"]
+__all__ = ["deprecated", "deprecated_parameter", "NeptuneDeprecationWarning"]
 
 
-def deprecated(*, alternative: Optional[str] = None):
+class NeptuneDeprecationWarning(DeprecationWarning):
+    pass
+
+
+warnings.simplefilter("always", category=NeptuneDeprecationWarning)
+
+
+warned_once = {}
+
+
+def warn_once(message: str, stack_level: int = 1):
+    if message not in warned_once.keys():
+        warnings.warn(
+            message=message,
+            category=NeptuneDeprecationWarning,
+            stacklevel=stack_level + 1,
+        )
+        warned_once[message] = 1
+
+
+def deprecated(*, alternative: Optional[str] = None, stack_level: int = 1):
     def deco(func):
         @wraps(func)
         def inner(*args, **kwargs):
             additional_info = f", use `{alternative}` instead" if alternative else " and will be removed"
 
-            warnings.simplefilter("once", DeprecationWarning)
-            warnings.warn(
-                f"`{func.__name__}` is deprecated{additional_info}."
+            warn_once(
+                message=f"`{func.__name__}` is deprecated{additional_info}."
                 f" We'll end support of it in `neptune-client==1.0.0`.",
-                category=DeprecationWarning,
-                stacklevel=2,
+                stack_level=stack_level + 1,
             )
-            warnings.simplefilter("default", DeprecationWarning)
 
             return func(*args, **kwargs)
 
@@ -52,14 +69,11 @@ def inner(*args, **kwargs):
                 if required_kwarg_name in kwargs:
                     raise NeptuneParametersCollision(required_kwarg_name, deprecated_kwarg_name, method_name=f.__name__)
 
-                warnings.simplefilter("once", DeprecationWarning)
-                warnings.warn(
-                    f"Parameter `{deprecated_kwarg_name}` is deprecated, use `{required_kwarg_name}` instead. We'll "
-                    f"end support of it in `neptune-client==1.0.0`.",
-                    category=DeprecationWarning,
-                    stacklevel=2,
+                warn_once(
+                    message=f"Parameter `{deprecated_kwarg_name}` is deprecated, use `{required_kwarg_name}` instead."
+                    " We'll end support of it in `neptune-client==1.0.0`.",
+                    stack_level=2,
                 )
-                warnings.simplefilter("default", DeprecationWarning)
 
                 kwargs[required_kwarg_name] = kwargs[deprecated_kwarg_name]
                 del kwargs[deprecated_kwarg_name]