Skip to content

Commit

Permalink
Added ability to determine when to do markup (#17)
Browse files Browse the repository at this point in the history
Adopts logic originally written for molecule that determines when
to use ANSI or not, relying on common environment variables.
  • Loading branch information
ssbarnea authored Nov 24, 2020
1 parent 138b6f1 commit 91d045a
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 3 deletions.
1 change: 1 addition & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ install_requires =
test =
mock>=3.0.5
pytest-cov>=2.7.1
pytest-mock>=3.3.1
pytest-plus
pytest-xdist>=1.29.0
pytest>=5.4.0
Expand Down
55 changes: 54 additions & 1 deletion src/enrich/console.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
"""Module that helps integrating with rich library."""
import io
import os
import sys
from typing import IO, Any, List, Union
from typing import IO, Any, List, TextIO, Union

import rich.console as rich_console
from rich.ansi import AnsiDecoder
Expand Down Expand Up @@ -62,6 +63,13 @@ def __init__(
self.redirect = redirect
self.soft_wrap = soft_wrap

# Unless user already mentioning terminal preference, we use our
# heuristic to make an informed decision.
if "force_terminal" not in kwargs:
kwargs["force_terminal"] = should_do_markup(
stream=kwargs.get("file", sys.stdout)
)

super().__init__(*args, **kwargs)
self.extended = True
if self.redirect:
Expand All @@ -85,3 +93,48 @@ def print(self, *args, **kwargs) -> None: # type: ignore
decoder = AnsiDecoder()
args = list(decoder.decode(text)) # type: ignore
super().print(*args, **kwargs)


# Based on Ansible implementation
def to_bool(value: Any) -> bool:
"""Return a bool for the arg."""
if value is None or isinstance(value, bool):
return bool(value)
if isinstance(value, str):
value = value.lower()
if value in ("yes", "on", "1", "true", 1):
return True
return False


def should_do_markup(stream: TextIO = sys.stdout) -> bool:
"""Decide about use of ANSI colors."""
py_colors = None

# https://xkcd.com/927/
for env_var in ["PY_COLORS", "CLICOLOR", "FORCE_COLOR", "ANSIBLE_FORCE_COLOR"]:
value = os.environ.get(env_var, None)
if value is not None:
py_colors = to_bool(value)
break

# If deliverately disabled colors
if os.environ.get("NO_COLOR", None):
return False

# User configuration requested colors
if py_colors is not None:
return to_bool(py_colors)

term = os.environ.get("TERM", "")
if "xterm" in term:
return True

if term == "dumb":
return False

# Use tty detection logic as last resort because there are numerous
# factors that can make isatty return a misleading value, including:
# - stdin.isatty() is the only one returning true, even on a real terminal
# - stderr returting false if user user uses a error stream coloring solution
return stream.isatty()
32 changes: 31 additions & 1 deletion src/enrich/test/test_console.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Tests for rich module."""
import sys

from enrich.console import Console
from enrich.console import Console, should_do_markup


def test_rich_console_ex() -> None:
Expand Down Expand Up @@ -52,5 +52,35 @@ def test_console_print_ansi() -> None:
assert "#00ff00" in html_result


def test_markup_detection_pycolors0(monkeypatch):
"""Assure PY_COLORS=0 disables markup."""
monkeypatch.setenv("PY_COLORS", "0")
assert not should_do_markup()


def test_markup_detection_pycolors1(monkeypatch):
"""Assure PY_COLORS=1 enables markup."""
monkeypatch.setenv("PY_COLORS", "1")
assert should_do_markup()


def test_markup_detection_tty_yes(mocker):
"""Assures TERM=xterm enables markup."""
mocker.patch("sys.stdout.isatty", return_value=True)
mocker.patch("os.environ", {"TERM": "xterm"})
assert should_do_markup()
mocker.resetall()
mocker.stopall()


def test_markup_detection_tty_no(mocker):
"""Assures that if no tty is reported we disable markup."""
mocker.patch("os.environ", {})
mocker.patch("sys.stdout.isatty", return_value=False)
assert not should_do_markup()
mocker.resetall()
mocker.stopall()


if __name__ == "__main__":
test_console_print_ansi()
2 changes: 1 addition & 1 deletion tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ passenv =
setenv =
PIP_DISABLE_VERSION_CHECK=1
PIP_USE_FEATURE={env:PIP_USE_FEATURE:2020-resolver}
PYTEST_REQPASS=5
PYTEST_REQPASS=9
PYTHONDONTWRITEBYTECODE=1
PYTHONUNBUFFERED=1
commands =
Expand Down

0 comments on commit 91d045a

Please sign in to comment.