From c6eb6e8b84562559aa332cdd167cc783bcaaa8b6 Mon Sep 17 00:00:00 2001
From: Maxime Armstrong <maxime@retain.ai>
Date: Fri, 27 Dec 2024 16:11:38 -0500
Subject: [PATCH] [utils] add beta warning and decorator

---
 .../dagster/dagster/_annotations.py           | 115 ++++++++++++++++++
 .../dagster/dagster/_utils/warnings.py        |  25 ++++
 2 files changed, 140 insertions(+)

diff --git a/python_modules/dagster/dagster/_annotations.py b/python_modules/dagster/dagster/_annotations.py
index 96d906aff0c3e..6966f4e208fed 100644
--- a/python_modules/dagster/dagster/_annotations.py
+++ b/python_modules/dagster/dagster/_annotations.py
@@ -11,6 +11,7 @@
     is_resource_def,
 )
 from dagster._utils.warnings import (
+    beta_warning,
     deprecation_warning,
     experimental_warning,
     preview_warning,
@@ -177,6 +178,120 @@ def get_preview_info(obj: Annotatable) -> PreviewInfo:
     return getattr(target, _PREVIEW_ATTR_NAME)
 
 
+# ########################
+# ##### BETA
+# ########################
+
+
+_BETA_ATTR_NAME: Final[str] = "_beta"
+
+
+@dataclass
+class BetaInfo:
+    additional_warn_text: Optional[str] = None
+    subject: Optional[str] = None
+
+
+@overload
+def beta(
+    __obj: T_Annotatable,
+    *,
+    additional_warn_text: Optional[str] = ...,
+    subject: Optional[str] = ...,
+    emit_runtime_warning: bool = ...,
+) -> T_Annotatable: ...
+
+
+@overload
+def beta(
+    __obj: None = ...,
+    *,
+    additional_warn_text: Optional[str] = ...,
+    subject: Optional[str] = ...,
+    emit_runtime_warning: bool = ...,
+) -> Callable[[T_Annotatable], T_Annotatable]: ...
+
+
+def beta(
+    __obj: Optional[T_Annotatable] = None,
+    *,
+    additional_warn_text: Optional[str] = None,
+    subject: Optional[str] = None,
+    emit_runtime_warning: bool = True,
+) -> Union[T_Annotatable, Callable[[T_Annotatable], T_Annotatable]]:
+    """Mark an object as beta. This appends some metadata to the object that causes it to be
+    rendered with a "beta" tag and associated warning in the docs.
+
+    If `emit_runtime_warning` is True, a warning will also be emitted when the function is called,
+    having the same text as is displayed in the docs. For consistency between docs and runtime
+    warnings, this decorator is preferred to manual calls to `beta_warning`.
+
+    Args:
+        additional_warn_text (Optional[str]): Additional text to display after the beta warning.
+        subject (Optional[str]): The subject of the beta warning. Defaults to a string
+            representation of the decorated object. This is useful when marking usage of
+            a beta API inside an otherwise non-beta function, so
+            that it can be easily cleaned up later. It should only be used with
+            `emit_runtime_warning=False`, as we don't want to warn users when a
+            beta API is used internally.
+        emit_runtime_warning (bool): Whether to emit a warning when the function is called.
+
+    Usage:
+
+        .. code-block:: python
+
+            @beta
+            def my_beta_function(my_arg):
+                ...
+
+            @beta
+            class MyBetaClass:
+                ...
+
+            @beta(subject="some_beta_function", emit_runtime_warning=False)
+            def not_beta_function():
+                ...
+                some_beta_function()
+                ...
+    """
+    if __obj is None:
+        return lambda obj: beta(
+            obj,
+            subject=subject,
+            emit_runtime_warning=emit_runtime_warning,
+            additional_warn_text=additional_warn_text,
+        )
+    else:
+        target = _get_annotation_target(__obj)
+        setattr(
+            target,
+            _BETA_ATTR_NAME,
+            BetaInfo(additional_warn_text, subject),
+        )
+
+        if emit_runtime_warning:
+            stack_level = _get_warning_stacklevel(__obj)
+            subject = subject or _get_subject(__obj)
+            warning_fn = lambda: beta_warning(
+                subject,
+                additional_warn_text=additional_warn_text,
+                stacklevel=stack_level,
+            )
+            return apply_pre_call_decorator(__obj, warning_fn)
+        else:
+            return __obj
+
+
+def is_beta(obj: Annotatable) -> bool:
+    target = _get_annotation_target(obj)
+    return hasattr(target, _BETA_ATTR_NAME)
+
+
+def get_beta_info(obj: Annotatable) -> BetaInfo:
+    target = _get_annotation_target(obj)
+    return getattr(target, _BETA_ATTR_NAME)
+
+
 # ########################
 # ##### SUPERSEDED
 # ########################
diff --git a/python_modules/dagster/dagster/_utils/warnings.py b/python_modules/dagster/dagster/_utils/warnings.py
index ac9e92f9d5fcc..135bab5bedfb3 100644
--- a/python_modules/dagster/dagster/_utils/warnings.py
+++ b/python_modules/dagster/dagster/_utils/warnings.py
@@ -35,6 +35,31 @@ def preview_warning(
     )
 
 
+# ########################
+# ##### BETA
+# ########################
+
+
+class BetaWarning(FutureWarning):
+    pass
+
+
+def beta_warning(
+    subject: str,
+    additional_warn_text: Optional[str] = None,
+    stacklevel: int = 3,
+):
+    if not _warnings_on.get():
+        return
+
+    warnings.warn(
+        f"{subject} is in beta stage - the feature is still being tested and may change."
+        + ((" " + additional_warn_text) if additional_warn_text else ""),
+        category=BetaWarning,
+        stacklevel=stacklevel,
+    )
+
+
 # ########################
 # ##### SUPERSEDED
 # ########################