add @checked (dagster-io#27134)

Probably about 6 years overdue, add a decorator that inserts a wrapping
function that performs `check` calls based on annotations.

## How I Tested These Changes

added test

Author: alangenfeld

python_modules/dagster/dagster/_check/decorator.py

@@ -0,0 +1,102 @@
+import inspect
+from abc import ABC
+from functools import update_wrapper
+from types import MethodType
+
+import dagster._check as check
+from dagster._check.builder import EvalContext, build_check_call_str
+
+
+class CheckedFnWrapper(ABC):
+    """Output of the @checked decorator, this class holds a reference to the decorated function
+    and upon first invocation compiles a wrapping function that performs run time type checks on
+    annotated inputs.
+
+    This class is not directly instantiated, but instead dynamic subclasses are created for each
+    callsite allowing the __call__ method to be replaced[1] with the compiled function to achieve
+    a single stack frame of over head for this decorator in best cases scenarios.
+
+
+    [1] __call__ can not be replaced on instances, only classes.
+    """
+
+    def __init__(self, fn):
+        self._target_fn = fn
+        self._eval_ctx = EvalContext.capture_from_frame(
+            2,
+            add_to_local_ns={},
+        )
+
+    def __get__(self, instance, _=None):
+        """Allow the decorated function to be bound to instances to support class methods."""
+        if instance:
+            return MethodType(self, instance)
+        return self
+
+    def __call__(self, *args, **kwargs):
+        signature = inspect.signature(self._target_fn)
+        lines = []
+        inputs = []
+
+        for name, param in signature.parameters.items():
+            if param.annotation != param.empty:
+                param_str = build_check_call_str(
+                    ttype=param.annotation,
+                    name=name,
+                    eval_ctx=self._eval_ctx,
+                )
+            else:
+                param_str = param.name
+
+            if param.kind in (param.KEYWORD_ONLY, param.POSITIONAL_OR_KEYWORD):
+                param_str = f"{param.name}={param_str}"
+            inputs.append(param.name)
+            lines.append(param_str)
+
+        lazy_imports_str = "\n    ".join(
+            f"from {module} import {t}" for t, module in self._eval_ctx.lazy_imports.items()
+        )
+
+        param_block = ",\n        ".join(lines)
+        inputs_block = ",\n    ".join(inputs)
+
+        checked_fn_name = f"__checked_{self._target_fn.__name__}"
+
+        fn_str = f"""
+def {checked_fn_name}(
+    __checked_wrapper,
+    {inputs_block}
+):
+    {lazy_imports_str}
+    return __checked_wrapper._target_fn(
+        {param_block}
+    )
+        """
+
+        if "check" not in self._eval_ctx.global_ns:
+            self._eval_ctx.global_ns["check"] = check
+
+        call = self._eval_ctx.compile_fn(
+            fn_str,
+            fn_name=checked_fn_name,
+        )
+
+        self.__class__.__call__ = call
+        return call(self, *args, **kwargs)
+
+
+def checked(fn):
+    """Decorator for adding runtime type checking based on type annotations."""
+    # if nothing can be checked, return the original fn
+    annotations = getattr(fn, "__annotations__", None)
+    if not annotations or (len(annotations) == 1 and set(annotations.keys()) == {"return"}):
+        return fn
+
+    # make a dynamic subclass to be able to hot swap __call__ post compilation
+    class _DynamicCheckedFnWrapper(CheckedFnWrapper): ...
+
+    checked_fn = _DynamicCheckedFnWrapper(fn)
+    return update_wrapper(
+        wrapper=checked_fn,
+        wrapped=fn,
+    )

python_modules/dagster/dagster_tests/general_tests/test_checked.py

@@ -0,0 +1,151 @@
+from typing import TYPE_CHECKING, Annotated, Optional
+
+import pytest
+from dagster._check import CheckError
+from dagster._check.builder import ImportFrom
+from dagster._check.decorator import checked
+
+if TYPE_CHECKING:
+    from dagster._core.test_utils import TestType
+
+
+def test_basic():
+    @checked
+    def foo(): ...
+
+    foo()
+
+    @checked
+    def bar(i: int): ...
+
+    bar(1)
+    bar(i=1)
+    with pytest.raises(CheckError):
+        bar("1")  # type: ignore
+    with pytest.raises(CheckError):
+        bar(i="1")  # type: ignore
+
+
+class Thing: ...
+
+
+def test_many():
+    @checked
+    def big(
+        name: str,
+        nick_names: list[str],
+        age: int,
+        cool: bool,
+        thing: Optional[Thing],
+        other_thing: Thing,
+        percent: float,
+        o_s: Optional[str],
+        o_n: Optional[int],
+        o_f: Optional[float],
+        o_b: Optional[bool],
+        foos: list[Annotated["TestType", ImportFrom("dagster._core.test_utils")]],
+    ):
+        return True
+
+    assert big(
+        name="dude",
+        nick_names=[
+            "foo",
+            "bar",
+            "biz",
+        ],
+        age=42,
+        cool=False,
+        thing=None,
+        other_thing=Thing(),
+        percent=0.5,
+        o_s="x",
+        o_n=3,
+        o_f=None,
+        o_b=None,
+        foos=[],
+    )
+
+    with pytest.raises(CheckError):
+        assert big(
+            name="dude",
+            nick_names=[
+                "foo",
+                "bar",
+                "biz",
+            ],
+            age=42,
+            cool=False,
+            thing=None,
+            other_thing=Thing(),
+            percent=0.5,
+            o_s="x",
+            o_n=3,
+            o_f="surprise_not_float",  # type: ignore
+            o_b=None,
+            foos=[],
+        )
+
+
+def test_no_op():
+    def foo(): ...
+
+    c_foo = checked(foo)
+    assert c_foo is foo
+
+    def bar() -> None: ...
+
+    c_bar = checked(bar)
+    assert c_bar is bar
+
+
+def test_star():
+    @checked
+    def baz(*, i: int): ...
+
+    baz(i=1)
+    with pytest.raises(CheckError):
+        baz(i="1")  # type: ignore
+
+
+def test_partial():
+    @checked
+    def foo(a, b, c: int): ...
+
+    foo(1, 2, 3)
+
+
+def test_class():
+    class Foo:
+        @checked
+        def me(self):
+            return self
+
+        @checked
+        def yell(self, word: str):
+            return word
+
+        @staticmethod
+        @checked
+        def holler(word: str):
+            return word
+
+        @classmethod
+        @checked
+        def scream(cls, word: str):
+            return word
+
+    f = Foo()
+    f.me()
+
+    f.yell("HI")
+    with pytest.raises(CheckError):
+        f.yell(3)  # type: ignore
+
+    Foo.holler("hi")
+    with pytest.raises(CheckError):
+        Foo.holler(3)  # type: ignore
+
+    Foo.scream("hi")
+    with pytest.raises(CheckError):
+        Foo.scream(3)  # type: ignore