Add AssetChecksDefinitionProps to replace untyped dictionary shenanigans

python_modules/dagster/dagster/_core/definitions/asset_checks.py

@@ -1,14 +1,16 @@
+from abc import ABC
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
+    Generic,
     Iterable,
     Iterator,
     Mapping,
     NamedTuple,
     Optional,
     Sequence,
     Set,
+    TypeVar,
 )
 
 from dagster import _check as check
@@ -29,15 +31,29 @@
 if TYPE_CHECKING:
     from dagster._core.definitions.assets import AssetsDefinition
 
+TPropsType = TypeVar("TPropsType", bound=NamedTuple)
+TDefType = TypeVar("TDefType")
 
-@experimental
-class AssetChecksDefinitionInputOutputProps(NamedTuple):
+
+class IDefinitionPropsClassMatchesInit(Generic[TPropsType, TDefType], ABC):
+    @staticmethod
+    def from_props(props: TPropsType) -> TDefType: ...
+
+
+class AssetChecksDefinitionProps(NamedTuple):
+    node_def: NodeDefinition
+    resource_defs: Mapping[str, ResourceDefinition]
+    specs: Sequence[AssetCheckSpec]
     asset_check_keys_by_output_name: Mapping[str, AssetCheckKey]
     asset_keys_by_input_name: Mapping[str, AssetKey]
 
 
 @experimental
-class AssetChecksDefinition(ResourceAddable, RequiresResources):
+class AssetChecksDefinition(
+    ResourceAddable,
+    RequiresResources,
+    IDefinitionPropsClassMatchesInit[AssetChecksDefinitionProps, "AssetChecksDefinition"],
+):
     """Defines a set of checks that are produced by the same op or op graph.
 
     AssetChecksDefinition are typically not instantiated directly, but rather produced using a
@@ -46,30 +62,52 @@ class AssetChecksDefinition(ResourceAddable, RequiresResources):
 
     def __init__(
         self,
-        *,
         node_def: NodeDefinition,
         resource_defs: Mapping[str, ResourceDefinition],
         specs: Sequence[AssetCheckSpec],
-        input_output_props: AssetChecksDefinitionInputOutputProps
-        # if adding new fields, make sure to handle them in the get_attributes_dict method
+        asset_check_keys_by_output_name: Mapping[str, AssetCheckKey],
+        asset_keys_by_input_name: Mapping[str, AssetKey],
     ):
-        self._node_def = node_def
-        self._resource_defs = resource_defs
-        self._specs = check.sequence_param(specs, "specs", of_type=AssetCheckSpec)
-        self._input_output_props = check.inst_param(
-            input_output_props, "input_output_props", AssetChecksDefinitionInputOutputProps
+        self._props = AssetChecksDefinitionProps(
+            node_def=check.inst_param(node_def, "node_def", NodeDefinition),
+            resource_defs=check.dict_param(
+                resource_defs, "resource_defs", key_type=str, value_type=ResourceDefinition
+            ),
+            specs=check.list_param(specs, "specs", of_type=AssetCheckSpec),
+            asset_check_keys_by_output_name=check.dict_param(
+                asset_check_keys_by_output_name,
+                "asset_check_keys_by_output_name",
+                key_type=str,
+                value_type=AssetCheckKey,
+            ),
+            asset_keys_by_input_name=check.dict_param(
+                asset_keys_by_input_name,
+                "asset_keys_by_input_name",
+                key_type=str,
+                value_type=AssetKey,
+            ),
         )
         self._specs_by_handle = {spec.key: spec for spec in specs}
         self._specs_by_output_name = {
             output_name: self._specs_by_handle[check_key]
-            for output_name, check_key in input_output_props.asset_check_keys_by_output_name.items()
+            for output_name, check_key in self._props.asset_check_keys_by_output_name.items()
         }
 
+    @staticmethod
+    def from_props(props: AssetChecksDefinitionProps) -> "AssetChecksDefinition":
+        return AssetChecksDefinition(
+            node_def=props.node_def,
+            resource_defs=props.resource_defs,
+            specs=props.specs,
+            asset_check_keys_by_output_name=props.asset_check_keys_by_output_name,
+            asset_keys_by_input_name=props.asset_keys_by_input_name,
+        )
+
     @public
     @property
     def node_def(self) -> NodeDefinition:
         """The op or op graph that can be executed to check the assets."""
-        return self._node_def
+        return self._props.node_def
 
     @public
     @property
@@ -109,11 +147,11 @@ def specs_by_output_name(self) -> Mapping[str, AssetCheckSpec]:
 
     @property
     def asset_keys_by_input_name(self) -> Mapping[str, AssetKey]:
-        return self._input_output_props.asset_keys_by_input_name
+        return self._props.asset_keys_by_input_name
 
     def get_resource_requirements(self) -> Iterator[ResourceRequirement]:
         yield from self.node_def.get_resource_requirements()  # type: ignore[attr-defined]
-        for source_key, resource_def in self._resource_defs.items():
+        for source_key, resource_def in self._props.resource_defs.items():
             yield from resource_def.get_resource_requirements(outer_context=source_key)
 
     def get_spec_for_check_key(self, asset_check_key: AssetCheckKey) -> AssetCheckSpec:
@@ -128,20 +166,14 @@ def required_resource_keys(self) -> Set[str]:
     def with_resources(
         self, resource_defs: Mapping[str, ResourceDefinition]
     ) -> "AssetChecksDefinition":
-        attributes_dict = self.get_attributes_dict()
-        attributes_dict["resource_defs"] = merge_resource_defs(
-            old_resource_defs=self._resource_defs,
-            resource_defs_to_merge_in=resource_defs,
-            requires_resources=self,
-        )
-        return self.__class__(**attributes_dict)
-
-    def get_attributes_dict(self) -> Dict[str, Any]:
-        return dict(
-            node_def=self._node_def,
-            resource_defs=self._resource_defs,
-            specs=self._specs,
-            input_output_props=self._input_output_props,
+        return self.__class__(
+            **self._props._replace(
+                resource_defs=merge_resource_defs(
+                    old_resource_defs=self._props.resource_defs,
+                    resource_defs_to_merge_in=resource_defs,
+                    requires_resources=self,
+                )
+            )._asdict()
         )
 
 

python_modules/dagster/dagster/_core/definitions/decorators/asset_check_decorator.py

@@ -8,7 +8,6 @@
 from dagster._core.definitions.asset_check_spec import AssetCheckSpec
 from dagster._core.definitions.asset_checks import (
     AssetChecksDefinition,
-    AssetChecksDefinitionInputOutputProps,
 )
 from dagster._core.definitions.assets import AssetsDefinition
 from dagster._core.definitions.events import AssetKey, CoercibleToAssetKey
@@ -170,12 +169,10 @@ def inner(fn: AssetCheckFunction) -> AssetChecksDefinition:
             node_def=op_def,
             resource_defs={},
             specs=[spec],
-            input_output_props=AssetChecksDefinitionInputOutputProps(
-                asset_keys_by_input_name={
-                    input_tuples_by_asset_key[resolved_asset_key][0]: resolved_asset_key
-                },
-                asset_check_keys_by_output_name={op_def.output_defs[0].name: spec.key},
-            ),
+            asset_keys_by_input_name={
+                input_tuples_by_asset_key[resolved_asset_key][0]: resolved_asset_key
+            },
+            asset_check_keys_by_output_name={op_def.output_defs[0].name: spec.key},
         )
 
         return checks_def

python_modules/dagster/dagster/_utils/test/__init__.py

@@ -3,6 +3,7 @@
 import tempfile
 from collections import defaultdict
 from contextlib import contextmanager
+from types import ModuleType
 from typing import Any, Dict, List, Mapping, Optional, Type, cast
 
 # top-level include is dangerous in terms of incurring circular deps
@@ -330,12 +331,20 @@ def check_concurrency_claim(
         return claim_status.with_sleep_interval(float(self._sleep_interval))
 
 
-def get_all_direct_subclasses_of_marker(marker_interface_cls: Type) -> List[Type]:
-    import dagster as dagster
+def get_all_direct_subclasses_of_marker(
+    marker_interface_cls: Type, module: Optional[ModuleType] = None
+) -> List[Type]:
+    """Get all direct subclasses of a given marker interface class from a given module. If the
+    module is not specified, it defaults to the dagster module.
+    """
+    if module is None:
+        import dagster as dagster
+
+        module = dagster
 
     return [
         symbol
-        for symbol in dagster.__dict__.values()
+        for symbol in module.__dict__.values()
         if isinstance(symbol, type)
         and issubclass(symbol, marker_interface_cls)
         and marker_interface_cls

python_modules/dagster/dagster_tests/definitions_tests/test_i_definition_prop_class_matches_init_adherence.py

@@ -0,0 +1,37 @@
+from inspect import signature
+from typing import List, Type
+
+from dagster._core.definitions import asset_checks
+from dagster._core.definitions.asset_checks import (
+    AssetChecksDefinition,
+    AssetChecksDefinitionProps,
+    IDefinitionPropsClassMatchesInit,
+)
+from dagster._utils.test import get_all_direct_subclasses_of_marker
+
+
+def test_adherence_of_all_idef_props_class_matches_init() -> None:
+    def_types: List[Type] = get_all_direct_subclasses_of_marker(
+        IDefinitionPropsClassMatchesInit, module=asset_checks
+    )
+
+    hard_coded_list = [
+        (AssetChecksDefinition, AssetChecksDefinitionProps),
+    ]
+
+    # keep hard_coded_list up-to-date with all subclasses of IDefinitionPropsClassMatchesInit
+    assert set([def_type.__name__ for def_type in def_types]) == set(
+        [entry[0].__name__ for entry in hard_coded_list]
+    ), (
+        "You likely added a subclass of IDefinitionPropsClassMatchesInit without adding it to the"
+        " hard_coded_list"
+    )
+
+    for def_type, props_type in hard_coded_list:
+        init_params = signature(def_type.__init__).parameters
+        init_param_names_minus_self = set(init_params.keys()) - set(["self"])
+        named_tuple_fields = set(props_type._fields)
+        assert init_param_names_minus_self == named_tuple_fields, (
+            f"You have added either added a field to {props_type.__name__} or removed a parameter"
+            f" from {def_type.__name__}.__init__. They must remain in sync"
+        )