Fix edge cases of generic @records (#24629)

## Summary & Motivation

This came about from the PR stacked above this. Realized that there were some remaining issues, particularly with `copy()` and type checking paramters that were generic. Added tests for these cases and fixed them.

## How I Tested These Changes

## Changelog

NOCHANGELOG

- [ ] `NEW` _(added new feature or capability)_
- [ ] `BUGFIX` _(fixed a bug)_
- [ ] `DOCS` _(added or updated documentation)_

Author: OwenKephart

python_modules/dagster/dagster/_check/builder.py

@@ -366,5 +366,11 @@ def build_check_call_str(
                 if tuple_types is not None:
                     tt_name = _name(tuple_types)
                     return f'{name} if isinstance({name}, {tt_name}) else check.inst_param({name}, "{name}", {tt_name})'
+        # generic
+        else:
+            inst_type = _coerce_type(ttype, eval_ctx)
+            if inst_type:
+                it = _name(inst_type)
+                return f'{name} if isinstance({name}, {it}) else check.inst_param({name}, "{name}", {it})'
 
         failed(f"Unhandled {ttype}")

python_modules/dagster/dagster/_record/__init__.py

@@ -128,6 +128,7 @@ def _namedtuple_record_transform(
         (cls, base),
         {  # these will override an implementation on the class if it exists
             **{n: getattr(base, n) for n in field_set.keys()},
+            "_fields": base._fields,
             "__iter__": _banned_iter,
             "__getitem__": _banned_idx,
             "__hidden_iter__": base.__iter__,
@@ -364,23 +365,27 @@ def __init__(
         self._nt_base = nt_base
         self._eval_ctx = eval_ctx
         self._new_frames = new_frames  # how many frames of __new__ there are
+        self._compiled_fn = None
 
     def __call__(self, cls, *args, **kwargs):
-        # update the context with callsite locals/globals to resolve
-        # ForwardRefs that were unavailable at definition time.
-        self._eval_ctx.update_from_frame(1 + self._new_frames)
-
-        # ensure check is in scope
-        if "check" not in self._eval_ctx.global_ns:
-            self._eval_ctx.global_ns["check"] = check
-
-        # jit that shit
-        self._nt_base.__new__ = self._eval_ctx.compile_fn(
-            self._build_checked_new_str(),
-            _CHECKED_NEW,
-        )
-
-        return self._nt_base.__new__(cls, *args, **kwargs)
+        if self._compiled_fn is None:
+            # update the context with callsite locals/globals to resolve
+            # ForwardRefs that were unavailable at definition time.
+            self._eval_ctx.update_from_frame(1 + self._new_frames)
+
+            # ensure check is in scope
+            if "check" not in self._eval_ctx.global_ns:
+                self._eval_ctx.global_ns["check"] = check
+
+            # we are double-memoizing this to handle some confusing mro issues
+            # in which the _nt_base's __new__ method is not on the critical
+            # path, causing this to get invoked multiple times
+            self._compiled_fn = self._eval_ctx.compile_fn(
+                self._build_checked_new_str(),
+                _CHECKED_NEW,
+            )
+            self._nt_base.__new__ = self._compiled_fn
+        return self._compiled_fn(cls, *args, **kwargs)
 
     def _build_checked_new_str(self) -> str:
         kw_args_str, set_calls_str = build_args_and_assignment_strs(self._field_set, self._defaults)

python_modules/dagster/dagster_tests/general_tests/test_record.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Sequence, TypeVar, Union
 
 import pytest
+from dagster._check.functions import CheckError
 from dagster._record import (
     _INJECTED_DEFAULT_VALS_LOCAL_VAR,
     IHaveNew,
@@ -600,3 +601,123 @@ class SubSub(Sub):
     assert subsub.b == 0
 
     assert repr(subsub) == "SubSub(a=0, b=0, c=-1, d=2)"
+
+
+def test_generic_with_propagate() -> None:
+    T = TypeVar("T")
+
+    class Base(Generic[T]): ...
+
+    @record
+    class RecordBase(Base[T]):
+        label: Optional[str] = None
+
+    @record
+    class SubAdditionalArg(RecordBase):
+        some_val: str
+
+    obj = SubAdditionalArg(some_val="hi")
+    assert SubAdditionalArg(some_val="hi").some_val == "hi"
+    assert copy(obj, label="...").label == "..."
+    assert copy(obj, some_val="new").some_val == "new"
+
+    @record
+    class SubAdditionalArgRecursive(RecordBase):
+        vals: Sequence[RecordBase]
+
+    obj = SubAdditionalArgRecursive(vals=[SubAdditionalArg(some_val="hi")])
+    assert len(obj.vals) == 1
+    assert copy(obj, label="...").label == "..."
+
+    @record
+    class SubAdditionalArgSpecific(RecordBase[int]):
+        vals: Sequence[RecordBase[str]]
+
+    obj = SubAdditionalArgSpecific(vals=[SubAdditionalArg(some_val="hi")])
+    assert len(obj.vals) == 1
+    assert copy(obj, label="...").label == "..."
+
+    @record
+    class SubAdditionalArgVariableBase(RecordBase[T]):
+        vals: Sequence[RecordBase[T]]
+        val: Base[T]
+
+    obj = SubAdditionalArgVariableBase(
+        vals=[SubAdditionalArg(some_val="hi")], val=SubAdditionalArg(some_val="bye")
+    )
+    assert len(obj.vals) == 1
+    assert copy(obj, label="...").label == "..."
+
+    @record
+    class SubSubVariableBaseSpecific(SubAdditionalArgVariableBase[str]): ...
+
+    obj = SubSubVariableBaseSpecific(
+        vals=[SubAdditionalArg(some_val="hi")], val=SubAdditionalArg(some_val="bye")
+    )
+    assert len(obj.vals) == 1
+    assert copy(obj, label="...").label == "..."
+
+    @record
+    class SubSubVariableBaseAny(SubAdditionalArgVariableBase): ...
+
+    obj = SubSubVariableBaseAny(
+        vals=[SubAdditionalArg(some_val="hi")], val=SubAdditionalArg(some_val="bye")
+    )
+    assert len(obj.vals) == 1
+    assert copy(obj, label="...").label == "..."
+
+
+def test_generic_with_propagate_type_checking() -> None:
+    T = TypeVar("T")
+
+    class Base(Generic[T]): ...
+
+    @record
+    class RecordBase(Base[T]):
+        inner: T
+
+    @record
+    class SpecificRecord(RecordBase):
+        val1: RecordBase
+        val2: RecordBase[str]
+        val3: Base
+        val4: Base[int]
+
+    valid_record = SpecificRecord(
+        inner=...,
+        val1=RecordBase(inner=1.23),
+        val2=RecordBase(inner="hi"),
+        val3=RecordBase(inner=1.23),
+        val4=RecordBase(inner=1),
+    )
+
+    with pytest.raises(CheckError, match='"val1" is not a RecordBase'):
+        copy(valid_record, val1=Base())
+
+    with pytest.raises(CheckError, match='"val2" is not a RecordBase'):
+        copy(valid_record, val2=Base())
+
+    with pytest.raises(CheckError, match='"val3" is not a Base'):
+        copy(valid_record, val3=3)
+
+    with pytest.raises(CheckError, match='"val4" is not a Base'):
+        copy(valid_record, val4=4)
+
+
+@pytest.mark.xfail()
+def test_custom_subclass() -> None:
+    @record_custom
+    class Thing(IHaveNew):
+        val: str
+
+        def __new__(cls, val_short: str):
+            return super().__new__(cls, val=val_short * 2)
+
+    assert Thing(val_short="abc").val == "abcabc"
+
+    @record
+    class SubThing(Thing):
+        other_val: int
+
+    # this does not work, as we've overridden the wrong __new__
+    SubThing(other_val=1)