allow flexible types for SourceAsset key arg (#7633)

dagster-io · Apr 28, 2022 · a83f387 · a83f387
1 parent cc07c17
commit a83f387
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/python_modules/dagster/dagster/core/asset_defs/source_asset.py b/python_modules/dagster/dagster/core/asset_defs/source_asset.py
@@ -1,7 +1,7 @@
 from typing import NamedTuple, Optional, Sequence, Union
 
 import dagster.check as check
-from dagster.core.definitions.events import AssetKey
+from dagster.core.definitions.events import AssetKey, CoerceableToAssetKey
 from dagster.core.definitions.metadata import (
     MetadataEntry,
     MetadataMapping,
@@ -28,7 +28,7 @@ class SourceAsset(
     that it's referenced from.
 
     Attributes:
-        key (AssetKey): The key of the asset.
+        key (Union[AssetKey, Sequence[str], str]): The key of the asset.
         metadata_entries (List[MetadataEntry]): Metadata associated with the asset.
         io_manager_key (str): The key for the IOManager that will be used to load the contents of
             the asset when it's used as an input to other assets inside a job.
@@ -39,7 +39,7 @@ class SourceAsset(
 
     def __new__(
         cls,
-        key: AssetKey,
+        key: CoerceableToAssetKey,
         metadata: Optional[MetadataUserInput] = None,
         io_manager_key: str = "io_manager",
         description: Optional[str] = None,
@@ -48,10 +48,9 @@ def __new__(
 
         metadata = check.opt_dict_param(metadata, "metadata", key_type=str)
         metadata_entries = normalize_metadata(metadata, [], allow_invalid=True)
-
         return super().__new__(
             cls,
-            key=check.inst_param(key, "key", AssetKey),
+            key=AssetKey.from_coerceable(key),
             metadata_entries=metadata_entries,
             io_manager_key=check.str_param(io_manager_key, "io_manager_key"),
             description=check.opt_str_param(description, "description"),

diff --git a/python_modules/dagster/dagster/core/definitions/events.py b/python_modules/dagster/dagster/core/definitions/events.py
@@ -146,6 +146,22 @@ def from_graphql_input(asset_key: Mapping[str, List[str]]) -> Optional["AssetKey
             return AssetKey(asset_key["path"])
         return None
 
+    @staticmethod
+    def from_coerceable(arg: "CoerceableToAssetKey") -> "AssetKey":
+        if isinstance(arg, AssetKey):
+            return check.inst_param(arg, "arg", AssetKey)
+        elif isinstance(arg, str):
+            return AssetKey([arg])
+        elif isinstance(arg, list):
+            check.list_param(arg, "arg", of_type=str)
+            return AssetKey(arg)
+        else:
+            check.tuple_param(arg, "arg", of_type=str)
+            return AssetKey(arg)
+
+
+CoerceableToAssetKey = Union[AssetKey, str, Sequence[str]]
+
 
 DynamicAssetKey = Callable[["OutputContext"], Optional[AssetKey]]
 
@@ -381,7 +397,7 @@ class AssetMaterialization(
 
     def __new__(
         cls,
-        asset_key: Union[List[str], AssetKey, str],
+        asset_key: CoerceableToAssetKey,
         description: Optional[str] = None,
         metadata_entries: Optional[List[Union[MetadataEntry, PartitionMetadataEntry]]] = None,
         partition: Optional[str] = None,

diff --git a/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_source_asset.py b/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_source_asset.py
@@ -17,3 +17,8 @@ def test_source_asset_metadata():
         "foo": MetadataValue.text("bar"),
         "baz": MetadataValue.text("[object] (unserializable)"),
     }
+
+
+def test_source_asset_key_args():
+    assert SourceAsset(key="foo").key == AssetKey(["foo"])
+    assert SourceAsset(key=["bar", "foo"]).key == AssetKey(["bar", "foo"])