Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
import io
import os
from collections import OrderedDict
from collections.abc import Generator
from .utils import DataClassType, LaxMapping, SQLAlchemyClassType, env_true, isatty
try:
from functools import cache
except ImportError:
from functools import lru_cache
cache = lru_cache()
__all__ = 'PrettyFormat', 'pformat', 'pprint'
MYPY = False
if MYPY:
from typing import Any, Iterable, Tuple, Union
PARENTHESES_LOOKUP = [
(list, '[', ']'),
(set, '{', '}'),
(frozenset, 'frozenset({', '})'),
]
DEFAULT_WIDTH = int(os.getenv('PY_DEVTOOLS_WIDTH', 120))
MISSING = object()
PRETTY_KEY = '__prettier_formatted_value__'
def fmt(v):
return {PRETTY_KEY: v}
class SkipPretty(Exception):
pass
@cache
def get_pygments():
try:
import pygments
from pygments.formatters import Terminal256Formatter
from pygments.lexers import PythonLexer
except ImportError: # pragma: no cover
return None, None, None
else:
return pygments, PythonLexer(), Terminal256Formatter(style='vim')
# common generator types (this is not exhaustive: things like chain are not include to avoid the import)
generator_types = Generator, map, filter, zip, enumerate
class PrettyFormat:
def __init__(
self,
indent_step=4,
indent_char=' ',
repr_strings=False,
simple_cutoff=10,
width=120,
yield_from_generators=True,
):
self._indent_step = indent_step
self._c = indent_char
self._repr_strings = repr_strings
self._repr_generators = not yield_from_generators
self._simple_cutoff = simple_cutoff
self._width = width
self._type_lookup = [
(dict, self._format_dict),
((str, bytes), self._format_str_bytes),
(tuple, self._format_tuples),
((list, set, frozenset), self._format_list_like),
(bytearray, self._format_bytearray),
(generator_types, self._format_generator),
# put this last as the check can be slow
(LaxMapping, self._format_dict),
(DataClassType, self._format_dataclass),
(SQLAlchemyClassType, self._format_sqlalchemy_class),
]
def __call__(self, value: 'Any', *, indent: int = 0, indent_first: bool = False, highlight: bool = False):
self._stream = io.StringIO()
self._format(value, indent_current=indent, indent_first=indent_first)
s = self._stream.getvalue()
pygments, pyg_lexer, pyg_formatter = get_pygments()
if highlight and pygments:
# apparently highlight adds a trailing new line we don't want
s = pygments.highlight(s, lexer=pyg_lexer, formatter=pyg_formatter).rstrip('\n')
return s
def _format(self, value: 'Any', indent_current: int, indent_first: bool):
if indent_first:
self._stream.write(indent_current * self._c)
try:
pretty_func = getattr(value, '__pretty__')
except AttributeError:
pass
else:
# `pretty_func.__class__.__name__ == 'method'` should only be true for bound methods,
# `hasattr(pretty_func, '__self__')` is more canonical but weirdly is true for unbound cython functions
from unittest.mock import _Call as MockCall
if pretty_func.__class__.__name__ == 'method' and not isinstance(value, MockCall):
try:
gen = pretty_func(fmt=fmt, skip_exc=SkipPretty)
self._render_pretty(gen, indent_current)
except SkipPretty:
pass
else:
return
value_repr = repr(value)
if len(value_repr) <= self._simple_cutoff and not isinstance(value, generator_types):
self._stream.write(value_repr)
else:
indent_new = indent_current + self._indent_step
for t, func in self._type_lookup:
if isinstance(value, t):
func(value, value_repr, indent_current, indent_new)
return
self._format_raw(value, value_repr, indent_current, indent_new)
def _render_pretty(self, gen, indent: int):
prefix = False
for v in gen:
if isinstance(v, int) and v in {-1, 0, 1}:
indent += v * self._indent_step
prefix = True
else:
if prefix:
self._stream.write('\n' + self._c * indent)
prefix = False
pretty_value = v.get(PRETTY_KEY, MISSING) if (isinstance(v, dict) and len(v) == 1) else MISSING
if pretty_value is not MISSING:
self._format(pretty_value, indent, False)
elif isinstance(v, str):
self._stream.write(v)
else:
# shouldn't happen but will
self._stream.write(repr(v))
def _format_dict(self, value: 'Any', _: str, indent_current: int, indent_new: int):
open_, before_, split_, after_, close_ = '{\n', indent_new * self._c, ': ', ',\n', '}'
if isinstance(value, OrderedDict):
open_, split_, after_, close_ = 'OrderedDict([\n', ', ', '),\n', '])'
before_ += '('
elif type(value) != dict:
open_, close_ = f'<{value.__class__.__name__}({{\n', '})>'
self._stream.write(open_)
for k, v in value.items():
self._stream.write(before_)
self._format(k, indent_new, False)
self._stream.write(split_)
self._format(v, indent_new, False)
self._stream.write(after_)
self._stream.write(indent_current * self._c + close_)
def _format_list_like(self, value: 'Union[list, tuple, set]', _: str, indent_current: int, indent_new: int):
open_, close_ = '(', ')'
for t, *oc in PARENTHESES_LOOKUP:
if isinstance(value, t):
open_, close_ = oc
break
self._stream.write(open_ + '\n')
for v in value:
self._format(v, indent_new, True)
self._stream.write(',\n')
self._stream.write(indent_current * self._c + close_)
def _format_tuples(self, value: tuple, value_repr: str, indent_current: int, indent_new: int):
fields = getattr(value, '_fields', None)
if fields:
# named tuple
self._format_fields(value, zip(fields, value), indent_current, indent_new)
else:
# normal tuples are just like other similar iterables
return self._format_list_like(value, value_repr, indent_current, indent_new)
def _format_str_bytes(self, value: 'Union[str, bytes]', value_repr: str, indent_current: int, indent_new: int):
if self._repr_strings:
self._stream.write(value_repr)
else:
lines = list(self._wrap_lines(value, indent_new))
if len(lines) > 1:
self._str_lines(lines, indent_current, indent_new)
else:
self._stream.write(value_repr)
def _str_lines(self, lines: 'Iterable[str]', indent_current: int, indent_new: int) -> None:
self._stream.write('(\n')
prefix = indent_new * self._c
for line in lines:
self._stream.write(prefix + repr(line) + '\n')
self._stream.write(indent_current * self._c + ')')
def _wrap_lines(self, s, indent_new) -> 'Generator[str, None, None]':
width = self._width - indent_new - 3
for line in s.splitlines(True):
start = 0
for pos in range(width, len(line), width):
yield line[start:pos]
start = pos
yield line[start:]
def _format_generator(self, value: Generator, value_repr: str, indent_current: int, indent_new: int):
if self._repr_generators:
self._stream.write(value_repr)
else:
name = value.__class__.__name__
if name == 'generator':
# no name if the name is just "generator"
self._stream.write('(\n')
else:
self._stream.write(f'{name}(\n')
for v in value:
self._format(v, indent_new, True)
self._stream.write(',\n')
self._stream.write(indent_current * self._c + ')')
def _format_bytearray(self, value: 'Any', _: str, indent_current: int, indent_new: int):
self._stream.write('bytearray')
lines = self._wrap_lines(bytes(value), indent_new)
self._str_lines(lines, indent_current, indent_new)
def _format_dataclass(self, value: 'Any', _: str, indent_current: int, indent_new: int):
from dataclasses import asdict
self._format_fields(value, asdict(value).items(), indent_current, indent_new)
def _format_sqlalchemy_class(self, value: 'Any', _: str, indent_current: int, indent_new: int):
fields = [
(field, getattr(value, field))
for field in dir(value)
if not (field.startswith('_') or field in ['metadata', 'registry'])
]
self._format_fields(value, fields, indent_current, indent_new)
def _format_raw(self, _: 'Any', value_repr: str, indent_current: int, indent_new: int):
lines = value_repr.splitlines(True)
if len(lines) > 1 or (len(value_repr) + indent_current) >= self._width:
self._stream.write('(\n')
wrap_at = self._width - indent_new
prefix = indent_new * self._c
from textwrap import wrap
for line in lines:
sub_lines = wrap(line, wrap_at)
for sline in sub_lines:
self._stream.write(prefix + sline + '\n')
self._stream.write(indent_current * self._c + ')')
else:
self._stream.write(value_repr)
def _format_fields(
self, value: 'Any', fields: 'Iterable[Tuple[str, Any]]', indent_current: int, indent_new: int
) -> None:
self._stream.write(f'{value.__class__.__name__}(\n')
for field, v in fields:
self._stream.write(indent_new * self._c)
if field: # field is falsy sometimes for odd things like call_args
self._stream.write(f'{field}=')
self._format(v, indent_new, False)
self._stream.write(',\n')
self._stream.write(indent_current * self._c + ')')
pformat = PrettyFormat()
force_highlight = env_true('PY_DEVTOOLS_HIGHLIGHT', None)
def pprint(s, file=None):
highlight = isatty(file) if force_highlight is None else force_highlight
print(pformat(s, highlight=highlight), file=file, flush=True)