Use AntlrAssetSelectionParser in AssetSelection.from_string

dagster-io · Nov 14, 2024 · 6e34fd3 · 6e34fd3
1 parent 456c239
commit 6e34fd3
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 21 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
@@ -495,27 +495,34 @@ 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
+        from dagster._core.definitions.antlr_asset_selection.antlr_asset_selection import (
+            AntlrAssetSelectionParser,
+        )
 
-        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}")