Skip to content

Commit

Permalink
Recreate AssetChecksDefinitions after setting prefix (#25271)
Browse files Browse the repository at this point in the history
## Summary & Motivation

This makes sure that the `load_asset_checks_from_*` functions return
`[AssetChecksDefinition]` by replacing the typecast with the creation of
an actual `AssetChecksDefinition`.

Alternatively, this issue could be addressed by adding a specialized
`with_attributes()` method to `AssetChecksDefinition` to make it return
the correct type. I'd be happy to try this path as well! But it feels
like this change would be sufficient and minimally impactful across the
codebase.

Closes #22064

## How I Tested These Changes

Added an `isinstance` check to the `load_asset_checks_from_*` tests
  • Loading branch information
marijncv authored Oct 15, 2024
1 parent ce0288b commit 3234086
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,18 @@ def _checks_with_attributes(
checks_defs: Sequence[AssetChecksDefinition],
asset_key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
) -> Sequence[AssetChecksDefinition]:
modified_checks = []
if asset_key_prefix:
modified_checks, _ = prefix_assets(checks_defs, asset_key_prefix, [], None)
return cast(Sequence[AssetChecksDefinition], modified_checks)
return [
AssetChecksDefinition.create(
keys_by_input_name=c.keys_by_input_name,
node_def=c.op,
check_specs_by_output_name=c.check_specs_by_output_name,
resource_defs=c.resource_defs,
can_subset=c.can_subset,
)
for c in modified_checks
]
else:
return checks_defs

Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import pytest
from dagster import (
AssetChecksDefinition,
AssetKey,
Definitions,
asset_check,
Expand All @@ -23,6 +24,7 @@ def test_load_asset_checks_from_modules():

checks = load_asset_checks_from_modules([checks_module])
assert len(checks) == 1
assert all(isinstance(check, AssetChecksDefinition) for check in checks)

asset_check_1_key = next(iter(asset_check_1.check_keys))

Expand All @@ -49,6 +51,7 @@ def test_load_asset_checks_from_modules_prefix():

checks = load_asset_checks_from_modules([checks_module], asset_key_prefix="foo")
assert len(checks) == 1
assert all(isinstance(check, AssetChecksDefinition) for check in checks)

check_key = next(iter(checks[0].check_keys))
assert check_key.asset_key == AssetKey(["foo", "asset_1"])
Expand Down Expand Up @@ -76,6 +79,7 @@ def check_in_current_module():
def test_load_asset_checks_from_current_module():
checks = load_asset_checks_from_current_module(asset_key_prefix="foo")
assert len(checks) == 1
assert all(isinstance(check, AssetChecksDefinition) for check in checks)
check_key = next(iter(checks[0].check_keys))
assert check_key.name == "check_in_current_module"
assert check_key.asset_key == AssetKey(["foo", "asset_1"])
Expand All @@ -100,6 +104,7 @@ def test_load_asset_checks_from_package(load_fns):

checks = checks_load_fn(checks_module, asset_key_prefix="foo")
assert len(checks) == 2
assert all(isinstance(check, AssetChecksDefinition) for check in checks)
check_key_0 = next(iter(checks[0].check_keys))
assert check_key_0.name == "asset_check_1"
assert check_key_0.asset_key == AssetKey(["foo", "asset_1"])
Expand Down

0 comments on commit 3234086

Please sign in to comment.