hack for output name

dagster-io · Jun 9, 2024 · 4161273 · 4161273
1 parent 12d17e4
commit 4161273
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/...dules/dagster/dagster/_core/definitions/decorators/decorator_assets_definition_builder.py b/...dules/dagster/dagster/_core/definitions/decorators/decorator_assets_definition_builder.py
@@ -44,7 +44,7 @@
 from dagster._core.types.dagster_type import DagsterType, Nothing
 
 from ..asset_check_spec import AssetCheckSpec
-from ..utils import NoValueSentinel
+from ..utils import DEFAULT_OUTPUT, NoValueSentinel
 from .op_decorator import _Op
 
 
@@ -465,7 +465,12 @@ def asset_keys(self) -> Set[AssetKey]:
 
     @cached_property
     def check_specs_by_output_name(self) -> Mapping[str, AssetCheckSpec]:
-        return create_check_specs_by_output_name(self.args.check_specs)
+        cs_by_name = create_check_specs_by_output_name(self.args.check_specs)
+        if self.args.decorator_name == "@asset_check":
+            check.invariant(len(cs_by_name) == 1, "Must have exactly one check spec")
+            return {DEFAULT_OUTPUT: next(iter(cs_by_name.values()))}
+
+        return cs_by_name
 
     @cached_property
     def check_outs_by_output_name(self) -> Mapping[str, Out]: