diff --git a/examples/development_to_production/development_to_production_tests/test_assets.py b/examples/development_to_production/development_to_production_tests/test_assets.py index fd5bc6b5d6699..19498e8cdb61f 100644 --- a/examples/development_to_production/development_to_production_tests/test_assets.py +++ b/examples/development_to_production/development_to_production_tests/test_assets.py @@ -1,12 +1,12 @@ import pandas as pd -from dagster import build_op_context +from dagster import build_asset_context from development_to_production.assets import comments, items, stories from development_to_production.resources import StubHNClient def test_items(): - context = build_op_context( + context = build_asset_context( resources={"hn_client": StubHNClient()}, op_config={"N": StubHNClient().fetch_max_item_id()} ) hn_dataset = items(context) diff --git a/examples/docs_snippets/docs_snippets/concepts/ops_jobs_graphs/unit_tests.py b/examples/docs_snippets/docs_snippets/concepts/ops_jobs_graphs/unit_tests.py index ba26159e4ae6e..4bf30b35fc959 100644 --- a/examples/docs_snippets/docs_snippets/concepts/ops_jobs_graphs/unit_tests.py +++ b/examples/docs_snippets/docs_snippets/concepts/ops_jobs_graphs/unit_tests.py @@ -11,6 +11,7 @@ Out, op, graph, + build_op_context, ) @@ -330,7 +331,7 @@ def test_asset_with_inputs(): # start_test_resource_asset -from dagster import asset, ConfigurableResource, build_op_context, with_resources +from dagster import asset, ConfigurableResource class BarResource(ConfigurableResource): @@ -351,7 +352,7 @@ def test_asset_requires_bar(): # start_test_config_asset -from dagster import asset, Config, build_op_context +from dagster import asset, Config class MyAssetConfig(Config): diff --git a/examples/docs_snippets/docs_snippets/guides/dagster/development_to_production/test_assets.py b/examples/docs_snippets/docs_snippets/guides/dagster/development_to_production/test_assets.py index 04794bde42a64..4391925cca5d4 100644 --- a/examples/docs_snippets/docs_snippets/guides/dagster/development_to_production/test_assets.py +++ b/examples/docs_snippets/docs_snippets/guides/dagster/development_to_production/test_assets.py @@ -1,6 +1,6 @@ import pandas as pd -from dagster import build_op_context +from dagster import build_asset_context from .assets_v2 import items from .resources.resources_v2 import StubHNClient @@ -10,7 +10,7 @@ def test_items(): - context = build_op_context( + context = build_asset_context( resources={"hn_client": StubHNClient()}, op_config={"N": StubHNClient().fetch_max_item_id()}, ) diff --git a/examples/docs_snippets/docs_snippets_tests/concepts_tests/resources_tests/test_pythonic_resources.py b/examples/docs_snippets/docs_snippets_tests/concepts_tests/resources_tests/test_pythonic_resources.py index 5533b7134cee6..d27168f42c55f 100644 --- a/examples/docs_snippets/docs_snippets_tests/concepts_tests/resources_tests/test_pythonic_resources.py +++ b/examples/docs_snippets/docs_snippets_tests/concepts_tests/resources_tests/test_pythonic_resources.py @@ -3,7 +3,6 @@ import mock import pytest -from dagster import build_init_resource_context, build_op_context from dagster._core.definitions.run_config import RunConfig from dagster._core.errors import DagsterInvalidConfigError diff --git a/python_modules/dagster/dagster_tests/core_tests/resource_tests/pythonic_resources/test_direct_invocation.py b/python_modules/dagster/dagster_tests/core_tests/resource_tests/pythonic_resources/test_direct_invocation.py index 24c1d30e5bb1e..e3ebb34c7206d 100644 --- a/python_modules/dagster/dagster_tests/core_tests/resource_tests/pythonic_resources/test_direct_invocation.py +++ b/python_modules/dagster/dagster_tests/core_tests/resource_tests/pythonic_resources/test_direct_invocation.py @@ -1,13 +1,9 @@ import pytest -from dagster import ( - ConfigurableResource, - asset, - op, -) +from dagster import AssetExecutionContext, ConfigurableResource, OpExecutionContext, asset, op from dagster._core.errors import ( DagsterInvalidInvocationError, ) -from dagster._core.execution.context.invocation import build_op_context +from dagster._core.execution.context.invocation import build_asset_context, build_op_context def test_direct_op_invocation() -> None: @@ -15,7 +11,7 @@ class MyResource(ConfigurableResource): a_str: str @op - def my_op(context, my_resource: MyResource) -> str: + def my_op(context: OpExecutionContext, my_resource: MyResource) -> str: assert my_resource.a_str == "foo" return my_resource.a_str @@ -61,7 +57,9 @@ class MyResource(ConfigurableResource): a_str: str @op - def my_op(context, my_resource: MyResource, my_other_resource: MyResource) -> str: + def my_op( + context: OpExecutionContext, my_resource: MyResource, my_other_resource: MyResource + ) -> str: assert my_resource.a_str == "foo" assert my_other_resource.a_str == "bar" return my_resource.a_str @@ -122,7 +120,9 @@ class MyResource(ConfigurableResource): z: int @op - def my_wacky_addition_op(context, x: int, y: int, my_resource: MyResource) -> int: + def my_wacky_addition_op( + context: OpExecutionContext, x: int, y: int, my_resource: MyResource + ) -> int: return x + y + my_resource.z # Just providing context is ok, we'll use the resource from the context @@ -192,12 +192,14 @@ class MyResource(ConfigurableResource): a_str: str @asset - def my_asset(context, my_resource: MyResource) -> str: + def my_asset(context: AssetExecutionContext, my_resource: MyResource) -> str: assert my_resource.a_str == "foo" return my_resource.a_str # Just providing context is ok, we'll use the resource from the context - assert my_asset(build_op_context(resources={"my_resource": MyResource(a_str="foo")})) == "foo" + assert ( + my_asset(build_asset_context(resources={"my_resource": MyResource(a_str="foo")})) == "foo" + ) # Providing both context and resource is not ok, because we don't know which one to use with pytest.raises( @@ -206,17 +208,17 @@ def my_asset(context, my_resource: MyResource) -> str: ): assert ( my_asset( - context=build_op_context(resources={"my_resource": MyResource(a_str="foo")}), + context=build_asset_context(resources={"my_resource": MyResource(a_str="foo")}), my_resource=MyResource(a_str="foo"), ) == "foo" ) # Providing resource only as kwarg is ok, we'll use that (we still need a context though) - assert my_asset(context=build_op_context(), my_resource=MyResource(a_str="foo")) == "foo" + assert my_asset(context=build_asset_context(), my_resource=MyResource(a_str="foo")) == "foo" # Providing resource as arg is ok, we'll use that (we still need a context though) - assert my_asset(build_op_context(), MyResource(a_str="foo")) == "foo" + assert my_asset(build_asset_context(), MyResource(a_str="foo")) == "foo" @asset def my_asset_no_context(my_resource: MyResource) -> str: @@ -225,7 +227,7 @@ def my_asset_no_context(my_resource: MyResource) -> str: # Providing context is ok, we just discard it and use the resource from the context assert ( - my_asset_no_context(build_op_context(resources={"my_resource": MyResource(a_str="foo")})) + my_asset_no_context(build_asset_context(resources={"my_resource": MyResource(a_str="foo")})) == "foo" ) @@ -238,28 +240,36 @@ class MyResource(ConfigurableResource): z: int @asset - def my_wacky_addition_asset(context, x: int, y: int, my_resource: MyResource) -> int: + def my_wacky_addition_asset( + context: AssetExecutionContext, x: int, y: int, my_resource: MyResource + ) -> int: return x + y + my_resource.z # Just providing context is ok, we'll use the resource from the context # We are successfully able to input x and y as args assert ( - my_wacky_addition_asset(build_op_context(resources={"my_resource": MyResource(z=2)}), 4, 5) + my_wacky_addition_asset( + build_asset_context(resources={"my_resource": MyResource(z=2)}), 4, 5 + ) == 11 ) # We can also input x and y as kwargs assert ( my_wacky_addition_asset( - build_op_context(resources={"my_resource": MyResource(z=3)}), y=1, x=2 + build_asset_context(resources={"my_resource": MyResource(z=3)}), y=1, x=2 ) == 6 ) # Providing resource only as kwarg is ok, we'll use that (we still need a context though) # We can input x and y as args - assert my_wacky_addition_asset(build_op_context(), 10, 20, my_resource=MyResource(z=30)) == 60 + assert ( + my_wacky_addition_asset(build_asset_context(), 10, 20, my_resource=MyResource(z=30)) == 60 + ) # We can also input x and y as kwargs in this case - assert my_wacky_addition_asset(build_op_context(), y=1, x=2, my_resource=MyResource(z=3)) == 6 + assert ( + my_wacky_addition_asset(build_asset_context(), y=1, x=2, my_resource=MyResource(z=3)) == 6 + ) @asset def my_wacky_addition_asset_no_context(x: int, y: int, my_resource: MyResource) -> int: @@ -269,14 +279,14 @@ def my_wacky_addition_asset_no_context(x: int, y: int, my_resource: MyResource) # We can input x and y as args assert ( my_wacky_addition_asset_no_context( - build_op_context(resources={"my_resource": MyResource(z=2)}), 4, 5 + build_asset_context(resources={"my_resource": MyResource(z=2)}), 4, 5 ) == 11 ) # We can also input x and y as kwargs assert ( my_wacky_addition_asset_no_context( - build_op_context(resources={"my_resource": MyResource(z=3)}), y=1, x=2 + build_asset_context(resources={"my_resource": MyResource(z=3)}), y=1, x=2 ) == 6 ) @@ -387,19 +397,21 @@ class NumResource(ConfigurableResource): executed = {} @asset - def an_asset(context, my_resource: NumResource, my_other_resource: NumResource) -> None: + def an_asset( + context: AssetExecutionContext, my_resource: NumResource, my_other_resource: NumResource + ) -> None: assert context.resources.my_resource.num == 1 assert context.resources.my_other_resource.num == 2 assert my_resource.num == 1 assert my_other_resource.num == 2 executed["yes"] = True - an_asset(build_op_context(), NumResource(num=1), NumResource(num=2)) + an_asset(build_asset_context(), NumResource(num=1), NumResource(num=2)) assert executed["yes"] executed.clear() an_asset( - build_op_context(), my_resource=NumResource(num=1), my_other_resource=NumResource(num=2) + build_asset_context(), my_resource=NumResource(num=1), my_other_resource=NumResource(num=2) ) assert executed["yes"] executed.clear() @@ -407,7 +419,7 @@ def an_asset(context, my_resource: NumResource, my_other_resource: NumResource) an_asset( my_other_resource=NumResource(num=2), my_resource=NumResource(num=1), - context=build_op_context(), + context=build_asset_context(), ) assert executed["yes"] executed.clear() diff --git a/python_modules/dagster/dagster_tests/core_tests/test_op_invocation.py b/python_modules/dagster/dagster_tests/core_tests/test_op_invocation.py index 8a52ae5417672..38099c6335e90 100644 --- a/python_modules/dagster/dagster_tests/core_tests/test_op_invocation.py +++ b/python_modules/dagster/dagster_tests/core_tests/test_op_invocation.py @@ -1240,7 +1240,7 @@ def my_asset(context): def non_asset_op(context): context.assets_def # noqa: B018 - with build_op_context( + with build_asset_context( partition_key="2023-02-02", ) as context: my_asset(context) @@ -1260,7 +1260,7 @@ def partitioned_asset(context): assert start == pendulum.instance(datetime(2023, 2, 2), tz=partitions_def.timezone) assert end == pendulum.instance(datetime(2023, 2, 3), tz=partitions_def.timezone) - context = build_op_context( + context = build_asset_context( partition_key="2023-02-02", ) partitioned_asset(context) @@ -1289,7 +1289,7 @@ def my_asset(context): assert context.asset_partitions_time_window_for_output() == time_window return 1 - context = build_op_context( + context = build_asset_context( partition_key="2020-01-01|a", ) my_asset(context) @@ -1309,7 +1309,7 @@ def static_multipartitioned_asset(context): ): context.asset_partitions_time_window_for_output() - context = build_op_context( + context = build_asset_context( partition_key="a|a", ) static_multipartitioned_asset(context) @@ -1323,7 +1323,7 @@ def foo(context: AssetExecutionContext): keys = partitions_def.get_partition_keys_in_range(context.partition_key_range) return {k: True for k in keys} - context = build_op_context( + context = build_asset_context( partition_key_range=PartitionKeyRange("2023-01-01", "2023-01-02"), ) assert foo(context) == {"2023-01-01": True, "2023-01-02": True}