import io

import os

from collections import OrderedDict

from collections.abc import Generator

from .utils import env_true, isatty

__all__ = 'PrettyFormat', 'pformat', 'pprint'

MYPY = False

if MYPY:

from typing import Any, 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

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')

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),

(Generator, self._format_generators),

]

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):

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

# very blunt check for things that look like dictionaries but do not necessarily inherit from Mapping

# e.g. asyncpg Records

# HELP: are there any other checks we should include here?

if (

hasattr(value, '__getitem__')

and hasattr(value, 'items')

and callable(value.items)

and not type(value) == type

):

self._format_dict(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', value_repr: 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 not isinstance(value, dict):

open_, close_ = '<{}({{\n'.format(value.__class__.__name__), '})>'

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]', value_repr: 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._stream.write(value.__class__.__name__ + '(\n')

for field, v in zip(fields, value):

self._stream.write(indent_new * self._c)

if field: # field is falsy sometimes for odd things like call_args

self._stream.write(str(field))

self._stream.write('=')

self._format(v, indent_new, False)

self._stream.write(',\n')

self._stream.write(indent_current * self._c + ')')

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._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 + ')')

else:

self._stream.write(value_repr)

def _wrap_lines(self, s, indent_new):

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_generators(self, value: Generator, value_repr: str, indent_current: int, indent_new: int):

if self._repr_generators:

self._stream.write(value_repr)

else:

self._stream.write('(\n')

for v in value:

self._format(v, indent_new, True)

self._stream.write(',\n')

self._stream.write(indent_current * self._c + ')')

def _format_raw(self, value: '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)

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)