avoid extra experimental warning when asset has partitions def (#7431)

dagster-io · Apr 13, 2022 · 6cf75eb · 6cf75eb
1 parent 7a9e8c9
commit 6cf75eb
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 9 deletions.
diff --git a/python_modules/dagster/dagster/core/asset_defs/decorators.py b/python_modules/dagster/dagster/core/asset_defs/decorators.py
@@ -184,16 +184,18 @@ def partition_fn(context):  # pylint: disable=function-redefined
                 return [context.partition_key]
 
         out_asset_key = AssetKey(list(filter(None, [*(self.namespace or []), asset_name])))
-        out = Out(
-            asset_key=out_asset_key,
-            metadata=self.metadata or {},
-            io_manager_key=self.io_manager_key,
-            dagster_type=self.dagster_type if self.dagster_type else NoValueSentinel,
-            asset_partitions_def=self.partitions_def,
-            asset_partitions=partition_fn,
-        )
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=ExperimentalWarning)
+
+            out = Out(
+                asset_key=out_asset_key,
+                metadata=self.metadata or {},
+                io_manager_key=self.io_manager_key,
+                dagster_type=self.dagster_type if self.dagster_type else NoValueSentinel,
+                asset_partitions_def=self.partitions_def,
+                asset_partitions=partition_fn,
+            )
+
             op = _Op(
                 name="__".join(out_asset_key.path),
                 description=self.description,

diff --git a/python_modules/dagster/dagster/core/definitions/output.py b/python_modules/dagster/dagster/core/definitions/output.py
@@ -401,7 +401,7 @@ def __new__(
         # make sure new parameters are updated in combine_with_inferred below
     ):
         if asset_partitions_def:
-            experimental_arg_warning("assets_definition", "Out.__new__")
+            experimental_arg_warning("asset_partitions_definition", "Out.__new__")
         return super(Out, cls).__new__(
             cls,
             dagster_type=NoValueSentinel

diff --git a/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_decorators.py b/python_modules/dagster/dagster_tests/core_tests/asset_defs_tests/test_decorators.py
@@ -8,6 +8,7 @@
     OpExecutionContext,
     Out,
     Output,
+    StaticPartitionsDefinition,
     String,
     build_op_context,
     check,
@@ -326,3 +327,13 @@ def asset_with_context(context, arg1):
     ctx = build_op_context()
     out = asset_with_context(ctx, 1)
     assert out == 1
+
+
+def test_partitions_def():
+    partitions_def = StaticPartitionsDefinition(["a", "b", "c", "d"])
+
+    @asset(partitions_def=partitions_def)
+    def my_asset():
+        pass
+
+    assert my_asset.partitions_def == partitions_def