From c4dfafeb261d6d9d2f8ffb1e8fb6ffc9b28f0faf Mon Sep 17 00:00:00 2001
From: Juliya Smith <jules@apeworx.io>
Date: Tue, 1 Aug 2023 11:48:41 -0500
Subject: [PATCH] fix: better cleanups

---
 src/ape/logging.py               | 14 +++++++++-----
 src/ape_networks/_cli.py         | 11 ++++++++++-
 tests/functional/test_logging.py |  7 +++++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/ape/logging.py b/src/ape/logging.py
index b717850221..c870cd025a 100644
--- a/src/ape/logging.py
+++ b/src/ape/logging.py
@@ -100,6 +100,7 @@ class CliLogger:
     def __init__(
         self,
         _logger: logging.Logger,
+        fmt: str,
         web3_request_logger: Optional[logging.Logger] = None,
         web3_http_logger: Optional[logging.Logger] = None,
     ):
@@ -112,6 +113,7 @@ def __init__(
         self._web3_request_manager_logger = web3_request_logger
         self._web3_http_provider_logger = web3_http_logger
         self._load_from_sys_argv()
+        self.fmt = fmt
 
     @classmethod
     def create(cls, fmt: Optional[str] = None, third_party: bool = True) -> "CliLogger":
@@ -121,14 +123,16 @@ def create(cls, fmt: Optional[str] = None, third_party: bool = True) -> "CliLogg
             kwargs["web3_http_logger"] = _get_logger("web3.providers.HTTPProvider", fmt=fmt)
 
         _logger = _get_logger("ape", fmt=fmt)
-        return cls(_logger, **kwargs)
+        return cls(_logger, fmt, **kwargs)
 
-    def format(self, new_format: str):
-        _format_logger(self._logger, new_format)
+    def format(self, fmt: Optional[str] = None):
+        self.fmt = fmt or DEFAULT_LOG_FORMAT
+        fmt = fmt or DEFAULT_LOG_FORMAT
+        _format_logger(self._logger, fmt)
         if req_log := self._web3_request_manager_logger:
-            _format_logger(req_log, new_format)
+            _format_logger(req_log, fmt)
         if prov_log := self._web3_http_provider_logger:
-            _format_logger(prov_log, new_format)
+            _format_logger(prov_log, fmt)
 
     def _load_from_sys_argv(self, default: Optional[Union[str, int]] = None):
         """
diff --git a/src/ape_networks/_cli.py b/src/ape_networks/_cli.py
index 7eb705017a..57c8a2aaeb 100644
--- a/src/ape_networks/_cli.py
+++ b/src/ape_networks/_cli.py
@@ -98,11 +98,20 @@ def run(cli_ctx, network):
         cli_ctx.abort("Process already running.")
 
     # Start showing process logs.
+    original_level = cli_ctx.logger.level
+    original_format = cli_ctx.logger.fmt
     cli_ctx.logger.set_level(LogLevel.DEBUG)
 
     # Change format to exclude log level (since it is always just DEBUG)
-    cli_ctx.logger.format("%(message)s")
+    cli_ctx.logger.format(fmt="%(message)s")
+    try:
+        _run(cli_ctx, provider)
+    finally:
+        cli_ctx.logger.set_level(original_level)
+        cli_ctx.logger.format(fmt=original_format)
+
 
+def _run(cli_ctx, provider: SubprocessProvider):
     provider.connect()
     if not (process := provider.process):
         provider.disconnect()
diff --git a/tests/functional/test_logging.py b/tests/functional/test_logging.py
index 44fa840b6e..4d78fea63f 100644
--- a/tests/functional/test_logging.py
+++ b/tests/functional/test_logging.py
@@ -90,8 +90,11 @@ def test_format(simple_runner):
     @group_for_testing.command()
     @ape_cli_context()
     def cmd(cli_ctx):
-        cli_ctx.logger.format("%(message)s")
-        cli_ctx.logger.success("this is a test")
+        cli_ctx.logger.format(fmt="%(message)s")
+        try:
+            cli_ctx.logger.success("this is a test")
+        finally:
+            cli_ctx.logger.format()
 
     result = simple_runner.invoke(group_for_testing, ["cmd", "-v", "SUCCESS"])
     assert "SUCCESS" not in result.output