Use AntlrAssetSelectionParser in AssetSelection.from_string

dagster-io · Nov 14, 2024 · 0615b89 · 0615b89
1 parent 197977a
commit 0615b89
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 22 deletions.
diff --git a/..._modules/dagster/dagster/_core/definitions/antlr_asset_selection/antlr_asset_selection.py b/..._modules/dagster/dagster/_core/definitions/antlr_asset_selection/antlr_asset_selection.py
@@ -150,5 +150,5 @@ def tree_str(self) -> str:
         return self._tree_str
 
     @property
-    def asset_selection(self) -> AssetSelection:
+    def asset_selection(self) -> "AssetSelection":
         return self._asset_selection
diff --git a/python_modules/dagster/dagster/_core/definitions/asset_selection.py b/python_modules/dagster/dagster/_core/definitions/asset_selection.py
@@ -8,6 +8,9 @@
 
 import dagster._check as check
 from dagster._annotations import deprecated, experimental, experimental_param, public
+from dagster._core.definitions.antlr_asset_selection.antlr_asset_selection import (
+    AntlrAssetSelectionParser,
+)
 from dagster._core.definitions.asset_check_spec import AssetCheckKey
 from dagster._core.definitions.asset_graph import AssetGraph
 from dagster._core.definitions.asset_key import (
@@ -495,27 +498,30 @@ def resolve_checks_inner(
 
     @classmethod
     def from_string(cls, string: str) -> "AssetSelection":
-        if string == "*":
-            return cls.all()
-
-        parts = parse_clause(string)
-        if parts is not None:
-            key_selection = cls.assets(parts.item_name)
-            if parts.up_depth and parts.down_depth:
-                selection = key_selection.upstream(parts.up_depth) | key_selection.downstream(
-                    parts.down_depth
-                )
-            elif parts.up_depth:
-                selection = key_selection.upstream(parts.up_depth)
-            elif parts.down_depth:
-                selection = key_selection.downstream(parts.down_depth)
-            else:
-                selection = key_selection
-            return selection
-
-        elif string.startswith("tag:"):
-            tag_str = string[len("tag:") :]
-            return cls.tag_string(tag_str)
+        try:
+            return AntlrAssetSelectionParser(string).asset_selection
+        except:
+            if string == "*":
+                return cls.all()
+
+            parts = parse_clause(string)
+            if parts is not None:
+                key_selection = cls.assets(parts.item_name)
+                if parts.up_depth and parts.down_depth:
+                    selection = key_selection.upstream(parts.up_depth) | key_selection.downstream(
+                        parts.down_depth
+                    )
+                elif parts.up_depth:
+                    selection = key_selection.upstream(parts.up_depth)
+                elif parts.down_depth:
+                    selection = key_selection.downstream(parts.down_depth)
+                else:
+                    selection = key_selection
+                return selection
+
+            elif string.startswith("tag:"):
+                tag_str = string[len("tag:") :]
+                return cls.tag_string(tag_str)
 
         check.failed(f"Invalid selection string: {string}")