Adds experimental_field decorator (#8066)

src/prefect/_internal/compatibility/experimental.py

@@ -12,12 +12,15 @@
 """
 import functools
 import warnings
-from typing import Any, Callable, Optional, Set, TypeVar
+from typing import Any, Callable, Optional, Set, Type, TypeVar
+
+import pydantic
 
 from prefect.settings import PREFECT_EXPERIMENTAL_WARN, SETTING_VARIABLES, Setting
 from prefect.utilities.callables import get_call_parameters
 
 T = TypeVar("T", bound=Callable)
+M = TypeVar("M", bound=pydantic.BaseModel)
 
 
 EXPERIMENTAL_WARNING = (
@@ -179,6 +182,62 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def experimental_field(
+    name: str,
+    *,
+    group: str,
+    help: str = "",
+    stacklevel: int = 2,
+    opt_in: bool = False,
+    when: Optional[Callable[[Any], bool]] = None,
+):
+    """
+    Mark a field in a Pydantic model as experimental.
+
+    Raises warning only if the field is specified during init.
+
+    Example:
+
+        ```python
+
+        @experimental_parameter("y", group="example", when=lambda y: y is not None)
+        def foo(x, y = None):
+            return x + 1 + (y or 0)
+        ```
+    """
+
+    when = when or (lambda _: True)
+
+    @experimental(
+        group=group,
+        feature=f"The field {name!r}",
+        help=help,
+        opt_in=opt_in,
+        stacklevel=stacklevel + 2,
+    )
+    def experimental_check():
+        """Utility function for performing a warning check for the specified group"""
+
+    # Replaces the model's __init__ method with one that performs an additional warning check
+    def decorator(model_cls: Type[M]) -> Type[M]:
+        cls_init = model_cls.__init__
+
+        @functools.wraps(model_cls.__init__)
+        def __init__(__pydantic_self__, **data: Any) -> None:
+            # Call the original init
+            cls_init(__pydantic_self__, **data)
+            # Perform warning check
+            if name in data.keys() and when(data[name]):
+                experimental_check()
+
+        # Patch the model's init method
+        model_cls.__init__ = __init__
+
+        return model_cls
+
+    return decorator
+
+
 def enabled_experiments() -> Set[str]:
     """
     Return the set of all enabled experiments.

tests/_internal/compatibility/test_experimental.py

@@ -1,13 +1,15 @@
 import re
 
 import pytest
+from pydantic import BaseModel, ValidationError
 
 from prefect._internal.compatibility.experimental import (
     ExperimentalFeature,
     ExperimentalFeatureDisabled,
     enabled_experiments,
     experiment_enabled,
     experimental,
+    experimental_field,
     experimental_parameter,
 )
 from prefect.settings import (
@@ -211,6 +213,92 @@ def foo(return_value: int = 1):
         foo(z=3)
 
 
+def test_experimental_field_warning():
+    @experimental_field(
+        "value",
+        group="test",
+        help="This is just a test, don't worry.",
+    )
+    class Foo(BaseModel):
+        value: int
+
+    with pytest.warns(
+        ExperimentalFeature,
+        match=(
+            "The field 'value' is experimental. This is just a test, "
+            "don't worry. The interface or behavior may change without warning, "
+            "we recommend pinning versions to prevent unexpected changes. "
+            "To disable warnings for this group of experiments, disable "
+            "PREFECT_EXPERIMENTAL_WARN_TEST."
+        ),
+    ):
+        assert Foo(value=2).value == 2
+
+
+def test_experimental_field_warning_no_warning_when_not_provided():
+    @experimental_field(
+        "value",
+        group="test",
+        help="This is just a test, don't worry.",
+    )
+    class Foo(BaseModel):
+        value: int = 1
+
+    assert Foo().value == 1
+
+
+def test_experimental_field_warning_when():
+    @experimental_field(
+        "value",
+        group="test",
+        help="This is just a test, don't worry.",
+        when=lambda x: x == 4,
+    )
+    class Foo(BaseModel):
+        value: int = 1
+
+    assert Foo(value=2).value == 2
+
+    with pytest.warns(
+        ExperimentalFeature,
+        match=(
+            "The field 'value' is experimental. This is just a test, "
+            "don't worry. The interface or behavior may change without warning, "
+            "we recommend pinning versions to prevent unexpected changes. "
+            "To disable warnings for this group of experiments, disable "
+            "PREFECT_EXPERIMENTAL_WARN_TEST."
+        ),
+    ):
+        assert Foo(value=4).value == 4
+
+
+def test_experimental_field_opt_in():
+    @experimental_field(
+        "value",
+        group="test",
+        help="This is just a test, don't worry.",
+        opt_in=True,
+    )
+    class Foo(BaseModel):
+        value: int = 1
+
+    with pytest.raises(ExperimentalFeatureDisabled):
+        assert Foo(value=1) == 1
+
+
+def test_experimental_field_retains_error_with_invalid_arguments():
+    @experimental_field(
+        "value",
+        group="test",
+        help="This is just a test, don't worry.",
+    )
+    class Foo(BaseModel):
+        value: int = 1
+
+    with pytest.raises(ValidationError, match="value is not a valid integer"):
+        Foo(value="nonsense")
+
+
 def test_experimental_warning_without_help():
     @experimental("A test function", group="test")
     def foo():