from __future__ import annotations
import contextlib
import enum
import os
import re
import typing
@typing.overload
def env_flag(name: str, default: bool) -> bool: ...
@typing.overload
def env_flag(name: str, default: bool | None = None) -> bool | None: ...
[docs]
def env_flag(name: str, default: bool | None = None) -> bool | None:
"""
Accepts environt variables formatted as y/n, yes/no, 1/0, true/false,
on/off, and returns it as a boolean.
If the environment variable is not defined, or has an unknown value,
returns `default`
"""
v = os.getenv(name)
if v and v.lower() in ('y', 'yes', 't', 'true', 'on', '1'):
return True
if v and v.lower() in ('n', 'no', 'f', 'false', 'off', '0'):
return False
return default
[docs]
class ColorSupport(enum.IntEnum):
"""Color support for the terminal."""
NONE = 0
XTERM = 16
XTERM_256 = 256
XTERM_TRUECOLOR = 16777216
WINDOWS = 8
[docs]
@classmethod
def from_env(cls) -> ColorSupport:
"""Get the color support from the environment.
If any of the environment variables contain `24bit` or `truecolor`,
we will enable true color/24 bit support. A `TERM` that is itself a
truecolor terminal (see `TRUECOLOR_TERMS`) also enables 24 bit
support. If they contain `256`, we will enable 256 color/8 bit
support. If they match a known ANSI terminal (see `ANSI_TERM_RE`,
e.g. `xterm-color`, `screen`, `tmux`, `konsole`, `rxvt`, `linux`), we
will enable 16 color support. Otherwise, we assume no color support.
If `JUPYTER_COLUMNS` or `JUPYTER_LINES` or `JPY_PARENT_PID` is set, we
will assume true color support.
Note that the highest available value will be used! Having
`COLORTERM=truecolor` will override `TERM=xterm-256color`.
"""
variables = (
'FORCE_COLOR',
'PROGRESSBAR_ENABLE_COLORS',
'COLORTERM',
'TERM',
)
# Precedence order is significant: an interactive Jupyter kernel and
# the Windows console probe each take priority over (and short-circuit)
# the env-var scan below.
if JUPYTER:
return cls._from_jupyter()
elif os.name == 'nt':
return cls._from_windows()
return cls._from_term_variables(variables)
@classmethod
def _from_jupyter(cls) -> ColorSupport:
"""Jupyter notebooks always support true color."""
return cls.XTERM_TRUECOLOR
@classmethod
def _from_windows(cls) -> ColorSupport: # pragma: no cover
"""Detect color support from the Windows console mode.
We can't reliably detect true color support on Windows, so we assume
it is supported when the console is configured to support it.
"""
from .terminal.os_specific import windows
if (
windows.get_console_mode()
& windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT
):
return cls.XTERM_TRUECOLOR
else:
return cls.WINDOWS
@classmethod
def _from_term_variables(
cls,
variables: tuple[str, ...],
) -> ColorSupport:
"""Pick the highest color support advertised by the terminal env vars.
The first `truecolor`/`24bit` value wins immediately; otherwise the
highest depth seen across all variables is returned. A generic truthy
flag such as `FORCE_COLOR=1` carries no depth and implies full color
support, analogous to the Jupyter handling above.
"""
support = cls.NONE
for variable in variables:
value = os.environ.get(variable)
if value is None:
continue
elif value in {'truecolor', '24bit'}:
# Truecolor support, we don't need to check anything else.
support = cls.XTERM_TRUECOLOR
break
elif value in TRUECOLOR_TERMS:
# A TERM name that itself guarantees a 24-bit terminal.
support = max(cls.XTERM_TRUECOLOR, support)
elif '256' in value:
support = max(cls.XTERM_256, support)
elif ANSI_TERM_RE.match(value):
# Any recognized ANSI terminal (xterm-color, screen, tmux,
# konsole, rxvt, linux, ...) advertises at least 16 colors,
# matching is_ansi_terminal()'s use of the same pattern.
support = max(cls.XTERM, support)
elif env_flag(variable, default=False):
return cls.XTERM_TRUECOLOR
return support
[docs]
def is_ansi_terminal(
fd: typing.IO[typing.Any],
is_terminal: bool | None = None,
) -> bool | None: # pragma: no cover
if is_terminal is None:
# Jupyter Notebooks support progress bars
if JUPYTER:
is_terminal = True
# This works for newer versions of pycharm only. With older versions
# there is no way to check.
elif os.environ.get('PYCHARM_HOSTED') == '1' and not os.environ.get(
'PYTEST_CURRENT_TEST'
):
is_terminal = True
if is_terminal is None:
# check if we are writing to a terminal or not. typically a file object
# is going to return False if the instance has been overridden and
# isatty has not been defined we have no way of knowing so we will not
# use ansi. ansi terminals will typically define one of the 2
# environment variables. Only the errors a stream legitimately
# produces are treated as "not a terminal": OSError (real I/O),
# ValueError (closed/detached file objects) and AttributeError
# (objects without isatty). Anything else is a bug and propagates.
with contextlib.suppress(OSError, ValueError, AttributeError):
is_tty: bool = fd.isatty()
# Try and match any of the huge amount of Linux/Unix ANSI consoles
if is_tty and ANSI_TERM_RE.match(os.environ.get('TERM', '')):
is_terminal = True
# ANSICON is a Windows ANSI compatible console
elif 'ANSICON' in os.environ:
is_terminal = True
elif os.name == 'nt':
from .terminal.os_specific import windows
return bool(
windows.get_console_mode()
& windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT,
)
else:
is_terminal = None
return is_terminal
[docs]
def is_terminal(
fd: typing.IO[typing.Any],
is_terminal: bool | None = None,
) -> bool | None:
if is_terminal is None:
# Full ansi support encompasses what we expect from a terminal
is_terminal = is_ansi_terminal(fd) or None
if is_terminal is None:
# Allow a environment variable override
is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None)
if is_terminal is None:
# If we do get a TTY we know this is a valid terminal. Streams can
# legitimately fail with OSError (real I/O), ValueError (closed or
# detached file objects) or AttributeError (no isatty at all);
# anything else is a bug and propagates.
try:
is_terminal = fd.isatty()
except (OSError, ValueError, AttributeError):
is_terminal = False
return is_terminal
JUPYTER = bool(
os.environ.get('JUPYTER_COLUMNS')
or os.environ.get('JUPYTER_LINES')
or os.environ.get('JPY_PARENT_PID')
)
ANSI_TERMS = (
'([xe]|bv)term',
'(sco)?ansi',
'cygwin',
'konsole',
'linux',
'rxvt',
'screen',
'tmux',
'vt(10[02]|220|320)',
)
ANSI_TERM_RE: re.Pattern[str] = re.compile(
f'^({"|".join(ANSI_TERMS)})', re.IGNORECASE
)
#: TERM values that on their own guarantee a truecolor-capable terminal, so
#: 24-bit color still engages when ``COLORTERM`` is stripped (e.g. over ssh
#: or sudo). Limited to names that *are* the terminal; generic values such as
#: ``xterm-256color`` are used by plenty of 256-only emulators.
TRUECOLOR_TERMS: frozenset[str] = frozenset({'xterm-kitty', 'xterm-ghostty'})
# Defined after ANSI_TERM_RE / TRUECOLOR_TERMS because from_env() reads them.
COLOR_SUPPORT = ColorSupport.from_env()