From 9d4b66786044989f85666a0e14d50ca3bed7b610 Mon Sep 17 00:00:00 2001 From: z80 <83730246+z80dev@users.noreply.github.com> Date: Sat, 18 Nov 2023 16:20:51 +0300 Subject: [PATCH] fix: run console under providercontext [APE-1529] (#1733) * fix: run console under providercontext * fix: actually persist state across exceptions * fix: add type hint to classvar * fix: lint * fix: add comments, recursive interactive check * fix: add test verifying state persists * fix: update test to use set_timestamp instead --- src/ape/api/networks.py | 33 +++++++++++++++++-- src/ape/cli/commands.py | 13 +++++++- src/ape/managers/networks.py | 7 +++- src/ape_run/_cli.py | 18 +++++----- .../cli/projects/script/scripts/error_cli.py | 4 +++ .../cli/projects/script/scripts/error_main.py | 5 +++ .../projects/script/scripts/error_no_def.py | 4 +++ tests/integration/cli/test_run.py | 3 +- 8 files changed, 73 insertions(+), 14 deletions(-) diff --git a/src/ape/api/networks.py b/src/ape/api/networks.py index c37386ffd4..1c37a73413 100644 --- a/src/ape/api/networks.py +++ b/src/ape/api/networks.py @@ -4,6 +4,7 @@ from typing import ( TYPE_CHECKING, Any, + ClassVar, Collection, Dict, Iterator, @@ -577,19 +578,45 @@ class ProviderContextManager(ManagerAccessMixin): provider_stack: List[str] = [] disconnect_map: Dict[str, bool] = {} - def __init__(self, provider: "ProviderAPI", disconnect_after: bool = False): + # We store a provider object at the class level for use when disconnecting + # due to an exception, when interactive mode is set. If we don't hold on + # to a reference to this object, the provider is dropped and reconnecting results + # in losing state when using a spawned local provider + _recycled_provider: ClassVar[Optional["ProviderAPI"]] = None + + def __init__( + self, + provider: "ProviderAPI", + disconnect_after: bool = False, + disconnect_on_exit: bool = True, + ): self._provider = provider self._disconnect_after = disconnect_after + self._disconnect_on_exit = disconnect_on_exit + self._skipped_disconnect = False @property def empty(self) -> bool: return not self.connected_providers or not self.provider_stack def __enter__(self, *args, **kwargs): + # If we have a recycled provider available, this means our last exit + # was due to an exception during interactive mode. We should resume that + # same connection, but also clear the object so we don't do this again + # in later provider contexts, which we would want to behave normally + if self._recycled_provider is not None: + # set inner var to the recycled provider for use in push_provider() + self._provider = self._recycled_provider + ProviderContextManager._recycled_provider = None return self.push_provider() - def __exit__(self, *args, **kwargs): - self.pop_provider() + def __exit__(self, exception, *args, **kwargs): + if not self._disconnect_on_exit and exception is not None: + # We want to skip disconnection when exiting due to an exception in interactive mode + if provider := self.network_manager.active_provider: + ProviderContextManager._recycled_provider = provider + else: + self.pop_provider() def push_provider(self): must_connect = not self._provider.is_connected diff --git a/src/ape/cli/commands.py b/src/ape/cli/commands.py index 8746b66294..13a3abbf6c 100644 --- a/src/ape/cli/commands.py +++ b/src/ape/cli/commands.py @@ -6,6 +6,16 @@ from ape import networks +def check_parents_for_interactive(ctx: Context) -> bool: + interactive: bool = ctx.params.get("interactive", False) + if interactive: + return True + # If not found, check the parent context. + if interactive is None and ctx.parent: + return check_parents_for_interactive(ctx.parent) + return False + + class NetworkBoundCommand(click.Command): """ A command that uses the :meth:`~ape.cli.options.network_option`. @@ -14,5 +24,6 @@ class NetworkBoundCommand(click.Command): def invoke(self, ctx: Context) -> Any: value = ctx.params.get("network") or networks.default_ecosystem.name - with networks.parse_network_choice(value): + interactive = check_parents_for_interactive(ctx) + with networks.parse_network_choice(value, disconnect_on_exit=not interactive): super().invoke(ctx) diff --git a/src/ape/managers/networks.py b/src/ape/managers/networks.py index 0458822447..6f3609b6d7 100644 --- a/src/ape/managers/networks.py +++ b/src/ape/managers/networks.py @@ -418,6 +418,7 @@ def parse_network_choice( network_choice: Optional[str] = None, provider_settings: Optional[Dict] = None, disconnect_after: bool = False, + disconnect_on_exit: bool = True, ) -> ProviderContextManager: """ Parse a network choice into a context manager for managing a temporary @@ -445,7 +446,11 @@ def parse_network_choice( provider = self.get_provider_from_choice( network_choice=network_choice, provider_settings=provider_settings ) - return ProviderContextManager(provider=provider, disconnect_after=disconnect_after) + return ProviderContextManager( + provider=provider, + disconnect_after=disconnect_after, + disconnect_on_exit=disconnect_on_exit, + ) @property def default_ecosystem(self) -> EcosystemAPI: diff --git a/src/ape_run/_cli.py b/src/ape_run/_cli.py index 737b6f9d8f..e9f180ce94 100644 --- a/src/ape_run/_cli.py +++ b/src/ape_run/_cli.py @@ -10,7 +10,7 @@ import click from click import Command, Context, Option -from ape import project +from ape import networks, project from ape.cli import NetworkBoundCommand, network_option, verbosity_option from ape.cli.options import _VERBOSITY_VALUES, _create_verbosity_kwargs from ape.exceptions import ApeException, handle_ape_exception @@ -76,13 +76,15 @@ def invoke(self, ctx: Context) -> Any: if ctx.params["interactive"]: # Print the exception trace and then launch the console # Attempt to use source-traceback style printing. - if not isinstance(err, ApeException) or not handle_ape_exception( - err, [ctx.obj.project_manager.path] - ): - err_info = traceback.format_exc() - click.echo(err_info) - - self._launch_console() + network_value = ctx.params.get("network") or networks.default_ecosystem.name + with networks.parse_network_choice(network_value, disconnect_on_exit=False): + if not isinstance(err, ApeException) or not handle_ape_exception( + err, [ctx.obj.project_manager.path] + ): + err_info = traceback.format_exc() + click.echo(err_info) + + self._launch_console() else: # Don't handle error - raise exception as normal. raise diff --git a/tests/integration/cli/projects/script/scripts/error_cli.py b/tests/integration/cli/projects/script/scripts/error_cli.py index 3403a28e8b..e7ad3793d2 100644 --- a/tests/integration/cli/projects/script/scripts/error_cli.py +++ b/tests/integration/cli/projects/script/scripts/error_cli.py @@ -1,7 +1,11 @@ import click +import ape + @click.command(short_help="Use a subcommand") def cli(): local_variable = "test foo bar" # noqa[F841] + provider = ape.chain.provider + provider.set_timestamp(123123123123123123) raise Exception("Expected exception") # noqa: T001 diff --git a/tests/integration/cli/projects/script/scripts/error_main.py b/tests/integration/cli/projects/script/scripts/error_main.py index 2f30e78f75..1f6cfbd208 100644 --- a/tests/integration/cli/projects/script/scripts/error_main.py +++ b/tests/integration/cli/projects/script/scripts/error_main.py @@ -1,3 +1,8 @@ +import ape + + def main(): local_variable = "test foo bar" # noqa[F841] + provider = ape.chain.provider + provider.set_timestamp(123123123123123123) raise Exception("Expected exception") diff --git a/tests/integration/cli/projects/script/scripts/error_no_def.py b/tests/integration/cli/projects/script/scripts/error_no_def.py index 1db3d6ab7a..7e8fae485d 100644 --- a/tests/integration/cli/projects/script/scripts/error_no_def.py +++ b/tests/integration/cli/projects/script/scripts/error_no_def.py @@ -1,2 +1,6 @@ +import ape + local_variable = "test foo bar" # noqa[F841] +provider = ape.chain.provider +provider.set_timestamp(123123123123123123) raise Exception("Expected exception") diff --git a/tests/integration/cli/test_run.py b/tests/integration/cli/test_run.py index b0faf306b4..78ed2340d0 100644 --- a/tests/integration/cli/test_run.py +++ b/tests/integration/cli/test_run.py @@ -100,13 +100,14 @@ def test_run_interactive(ape_cli, runner, project): ] # Show that the variable namespace from the script is available in the console. - user_input = "local_variable\nexit\n" + user_input = "local_variable\nape.chain.provider.mine()\nape.chain.blocks.head\nexit\n" result = runner.invoke(ape_cli, ["run", "--interactive", scripts[0].stem], input=user_input) assert result.exit_code == 0, result.output # From script: local_variable = "test foo bar" assert "test foo bar" in result.output + assert "timestamp=123123123123123" in result.output @skip_projects_except("script")