Skip to content

Commit

Permalink
perf: logger module loading (#2408)
Browse files Browse the repository at this point in the history
  • Loading branch information
antazoey authored Dec 11, 2024
1 parent f924acf commit fbdbe52
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 12 deletions.
29 changes: 18 additions & 11 deletions src/ape/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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]
Expand All @@ -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.
Expand All @@ -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"]
11 changes: 10 additions & 1 deletion tests/functional/test_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:[email protected]/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

0 comments on commit fbdbe52

Please sign in to comment.