From fbdbe52d286d12abe400577c3e7e670ee8a947f8 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Thu, 12 Dec 2024 03:28:21 +0700
Subject: [PATCH] perf: logger module loading (#2408)

---
 src/ape/logging.py               | 29 ++++++++++++++++++-----------
 tests/functional/test_logging.py | 11 ++++++++++-
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/ape/logging.py b/src/ape/logging.py
index 34fa511e1e..e3c88960d5 100644
--- a/src/ape/logging.py
+++ b/src/ape/logging.py
@@ -6,11 +6,13 @@
 from contextlib import contextmanager
 from enum import IntEnum
 from pathlib import Path
-from typing import IO, Any, Optional, Union
+from typing import IO, TYPE_CHECKING, Any, Optional, Union
+from urllib.parse import urlparse, urlunparse
 
 import click
-from rich.console import Console as RichConsole
-from yarl import URL
+
+if TYPE_CHECKING:
+    from rich.console import Console as RichConsole
 
 
 class LogLevel(IntEnum):
@@ -325,23 +327,28 @@ def _get_level(level: Optional[Union[str, int, LogLevel]] = None) -> str:
 def sanitize_url(url: str) -> str:
     """Removes sensitive information from given URL"""
 
-    url_obj = URL(url).with_user(None).with_password(None)
+    parsed = urlparse(url)
+    new_netloc = parsed.hostname or ""
+    if parsed.port:
+        new_netloc += f":{parsed.port}"
 
-    # If there is a path, hide it but show that you are hiding it.
-    # Use string interpolation to prevent URL-character encoding.
-    return f"{url_obj.with_path('')}/{HIDDEN_MESSAGE}" if url_obj.path else f"{url}"
+    new_url = urlunparse(parsed._replace(netloc=new_netloc, path=""))
+    return f"{new_url}/{HIDDEN_MESSAGE}" if parsed.path else new_url
 
 
 logger = ApeLogger.create()
 
 
 class _RichConsoleFactory:
-    rich_console_map: dict[str, RichConsole] = {}
+    rich_console_map: dict[str, "RichConsole"] = {}
 
-    def get_console(self, file: Optional[IO[str]] = None, **kwargs) -> RichConsole:
+    def get_console(self, file: Optional[IO[str]] = None, **kwargs) -> "RichConsole":
         # Configure custom file console
         file_id = str(file)
         if file_id not in self.rich_console_map:
+            # perf: delay importing from rich, as it is slow.
+            from rich.console import Console as RichConsole
+
             self.rich_console_map[file_id] = RichConsole(file=file, width=100, **kwargs)
 
         return self.rich_console_map[file_id]
@@ -350,7 +357,7 @@ def get_console(self, file: Optional[IO[str]] = None, **kwargs) -> RichConsole:
 _factory = _RichConsoleFactory()
 
 
-def get_rich_console(file: Optional[IO[str]] = None, **kwargs) -> RichConsole:
+def get_rich_console(file: Optional[IO[str]] = None, **kwargs) -> "RichConsole":
     """
     Get an Ape-configured rich console.
 
@@ -361,7 +368,7 @@ def get_rich_console(file: Optional[IO[str]] = None, **kwargs) -> RichConsole:
     Returns:
         ``rich.Console``.
     """
-    return _factory.get_console(file)
+    return _factory.get_console(file, **kwargs)
 
 
 __all__ = ["DEFAULT_LOG_LEVEL", "logger", "LogLevel", "ApeLogger", "get_rich_console"]
diff --git a/tests/functional/test_logging.py b/tests/functional/test_logging.py
index 72bd0dd313..bee67db477 100644
--- a/tests/functional/test_logging.py
+++ b/tests/functional/test_logging.py
@@ -3,7 +3,7 @@
 from click.testing import CliRunner
 
 from ape.cli import ape_cli_context
-from ape.logging import LogLevel, logger
+from ape.logging import LogLevel, logger, sanitize_url
 
 
 @pytest.fixture
@@ -121,3 +121,12 @@ def test_at_level():
         assert logger.level == level_to_set
 
     assert logger.level == initial_level
+
+
+@pytest.mark.parametrize(
+    "url", ("https://user:password@example.com/v1/API_KEY", "https://example.com/v1/API_KEY")
+)
+def test_sanitize_url(url):
+    actual = sanitize_url(url)
+    expected = "https://example.com/[hidden]"
+    assert actual == expected