AssetSelection.assets (#8316)

dagster-io · Jun 10, 2022 · e6fad5b · e6fad5b
1 parent 2623403
commit e6fad5b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/python_modules/dagster/dagster/core/asset_defs/asset_selection.py b/python_modules/dagster/dagster/core/asset_defs/asset_selection.py
@@ -18,6 +18,12 @@ class AssetSelection(ABC):
     def all() -> "AllAssetSelection":
         return AllAssetSelection()
 
+    @staticmethod
+    def assets(*assets_defs: AssetsDefinition) -> "KeysAssetSelection":
+        return KeysAssetSelection(
+            *(key for assets_def in assets_defs for key in assets_def.asset_keys)
+        )
+
     @staticmethod
     def keys(*asset_keys: CoerceableToAssetKey) -> "KeysAssetSelection":
         _asset_keys = [AssetKey.from_coerceable(key) for key in asset_keys]

diff --git a/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_asset_selection.py b/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_asset_selection.py
@@ -94,6 +94,11 @@ def test_asset_selection_keys(all_assets):
     assert sel.resolve(all_assets) == _asset_keys_of({alice, bob})
 
 
+def test_asset_selection_assets(all_assets):
+    sel = AssetSelection.assets(alice, bob)
+    assert sel.resolve(all_assets) == _asset_keys_of({alice, bob})
+
+
 def test_asset_selection_or(all_assets):
     sel = AssetSelection.keys("alice", "bob") | AssetSelection.keys("bob", "candace")
     assert sel.resolve(all_assets) == _asset_keys_of({alice, bob, candace})