"""
Make assert statements print pretty output, including source.
This makes ``assert x == y`` more usable, as an alternative to
``self.assertEqual(x, y)``
This plugin implements :func:`outcomeDetail` and checks for event.exc_info
If it finds that an AssertionError happened, it will inspect the traceback and
add additional detail to the error report.
"""
import collections
import inspect
import io
import re
import textwrap
import tokenize
from nose2 import events
__unittest = True
[文档]
class PrettyAssert(events.Plugin):
"""Add pretty output for "assert" statements"""
configSection = "pretty-assert"
commandLineSwitch = (
None,
"pretty-assert",
'Add pretty output for "assert" statements',
)
def outcomeDetail(self, event):
# skip if no exception or expected error
if (not event.outcomeEvent.exc_info) or event.outcomeEvent.expected:
return
# unpack, but skip if it's not an AssertionError
excty, exc, trace = event.outcomeEvent.exc_info
if excty is not AssertionError:
return
self.addAssertDetail(event.extraDetail, exc, trace)
[文档]
@staticmethod
def addAssertDetail(extraDetail, exc, trace):
"""
Add details to output regarding AssertionError and its context
extraDetail: a list of lines which will be joined with newlines and
added to the output for this test failure -- defined as part of the
event format
exc: the AssertionError exception which was thrown
trace: a traceback object for the exception
"""
assert_statement, token_descriptions = _collect_assert_data(trace)
# no message was given
if len(exc.args) == 0:
message = None
else:
message = exc.args[0]
# if there is no assertion statement found, do not add detail to output
#
# in cases like unittest assert*() methods, an assertion error is
# raised, but it doesn't originate with an `assert` statement and has
# an autogenerated message
if not assert_statement:
return
#
# actually add exception info to detail
#
# add the assert statement to output with '>>>' prefix
extraDetail.append(re.sub("^", ">>> ", assert_statement, flags=re.MULTILINE))
if message:
extraDetail.append("\nmessage:")
extraDetail.append(f" {message}")
if token_descriptions:
extraDetail.append("\nvalues:")
for k, v in token_descriptions.items():
extraDetail.append(f" {k} = {v}")
def _collect_assert_data(trace):
"""
Given a traceback, extract the assertion statement and get the set of bound
variable names (i.e. tokens)
"""
# inspect the trace, collecting various data and determining whether or not
# it can be tokenized at all
source_lines, frame_locals, frame_globals, can_tokenize = _get_inspection_info(
trace
)
# if things will tokenize cleanly, actually do it
if can_tokenize:
assert_startline, token_descriptions = _tokenize_assert(
source_lines, frame_locals, frame_globals
)
# otherwise, indicate that we can't render detail by use of Nones
else:
assert_startline = None
token_descriptions = None
# if we found an "assert" (we might not, if someone raises
# AssertionError themselves), grab the whole assertion statement
#
# as a fallback, stick with whatever we think the statement was
# - this is easily deceived by multiline expressions
if assert_startline is not None:
statement = textwrap.dedent(
"".join(source_lines[assert_startline:]).rstrip("\n")
)
else:
statement = None
return statement, token_descriptions
def _get_inspection_info(trace):
"""
Pick apart a traceback for the info we actually want to inspect from it
- lines of source (truncated)
- locals and globals from execution frame
- statement which failed (which can be garbage -- don't trust it)
- can_tokenize: a bool indicating that the lines of source can be parsed
"""
(frame, fname, lineno, funcname, context, ctx_index) = inspect.getinnerframes(
trace
)[-1]
original_source_lines, firstlineno = inspect.getsourcelines(frame)
# truncate to the code in this frame to remove anything after current
# assert statement
last_index = lineno - firstlineno + 1
source_lines = original_source_lines[:last_index]
# in case the current line is actually an incomplete expression, as in
# assert x == (y
# ).z
#
# in which case the the current line is "assert x == (y", which is not a
# complete expression
# try to append lines to complete the expression, retrying parsing each and
# every time until it succeeds
for line in original_source_lines[last_index:]:
if _can_tokenize(source_lines):
break
else:
source_lines.append(line)
return (source_lines, frame.f_locals, frame.f_globals, _can_tokenize(source_lines))
def _can_tokenize(source_lines):
"""
Check if a list of lines of source can successfully be tokenized
"""
# tokenize.generate_tokens requires a file-like object, so we need to
# convert source_lines to a StringIO to give it that interface
filelike = io.StringIO(textwrap.dedent("".join(source_lines)))
try:
for _tokty, _tok, _start, _end, _tok_lineno in tokenize.generate_tokens(
filelike.readline
):
pass
except tokenize.TokenError:
return False
return True
def _tokenize_assert(source_lines, frame_locals, frame_globals):
"""
Given a set of lines of source ending in a failing assert, plus the frame
locals and globals, tokenize source.
Only look at tokens in the final assert statement
Resolve all names to repr() of values
Return
The line on which the assert starts (relative to start of
source_lines)
A collection of token descriptions as a name=val ordered dict
"""
# tokenize.generate_tokens requires a file-like object, so we need to
# convert source_lines to a StringIO to give it that interface
filelike_context = io.StringIO(textwrap.dedent("".join(source_lines)))
# track the first line of the assert statement
# when the assert is on oneline, we'll have it easily, but a multiline
# statement like
# assert (x ==
# 1)
# will leave us holding the last line of the statement,
# e.g. " 1)", which is not useful
# so every time a new assert is found, we get a value back indicate
# that it's the start line
#
# assert True
# assert False
# works fine, because we'll just hold the last value
#
# assert True
# assert False
# assert True
# also works because we truncated source_lines to remove the final
# assert, which we didn't reach during execution
assert_startline = None
token_processor = TokenProcessor(frame_locals, frame_globals)
# tokenize and process each token
for tokty, tok, start, end, tok_lineno in tokenize.generate_tokens(
filelike_context.readline
):
ret = token_processor.handle_token(tokty, tok, start, end, tok_lineno)
if ret:
assert_startline = ret
# adjust assert_startline by 1 to become a valid index into the
# source_lines -- "line 1" means "index 0"
if assert_startline:
assert_startline -= 1
token_descriptions = collections.OrderedDict()
for name, obj in token_processor.get_token_collection().items():
# okay, get repr() for a good string representation
strvalue = repr(obj)
# add in the form we want to print
token_descriptions[name] = strvalue
return assert_startline, token_descriptions
class TokenProcessor:
def __init__(self, frame_locals, frame_globals):
# local and global variables from the frame which we're inspecting
self.frame_locals, self.frame_globals = frame_locals, frame_globals
# None or a tuple of (object, name) where
# - "object" is the object whose attributes we are currently resolving
# - "name" is its name, as we would like to display it
#
# start each time we see a sequence of NAME OP NAME OP NAME (etc.)
# end each time we see a token which is neither NAME nor OP
self.doing_resolution = None
# an index of known token names (including the long "x.y.z" names we
# get from attribute resolution) to their values, in the order in which
# they were encountered
# track which tokens we've seen to avoid duplicates if a name appears
# twice, as in `assert x != x`
self.seen_tokens = collections.OrderedDict()
# the previous token seen as a tuple of (tok_type, token_name)
# (or None when we start)
self.last_tok = None
def get_token_collection(self):
return self.seen_tokens
def handle_token(self, toktype, tok, start, end, line):
"""
A tokenization processor for tokenize.generate_tokens
Skips certain token types, class names, etc
When an identifiable/usable token is found, add it to the token
collection (self.seen_tokens)
When an "assert" statement is found, reset the token collection
and return the start line (relative to the text being tokenized)
"""
prior_tok = self.last_tok
self.last_tok = (toktype, tok)
# CASE 0: skip non "NAME" or "OP" tokens and clear current resolution
#
# NAME is most identifiers and keywords
# OP is operators, including .
#
# special note: don't clear resolution for whitespace (e.g. newline)
if toktype not in (tokenize.NAME, tokenize.OP):
# only newline for now, maybe we'll find others
if toktype not in (tokenize.NEWLINE,):
self.doing_resolution = None
return
# CASE 1: Operator token
#
# skip tokens and either leave resolution in progress or reset,
# depending
#
# continue resolution for
# "."
# because that's what attribute resolution *is*
# ")"
# this is handy, as it means that "(x).y" works
#
# reset resolution for everything else, e.g. "[", "]", ":"
# special note: reset resolution for "("
#
# failing to filter out "(" can result in badness in cases like this:
# >>> def foo():
# >>> return [1]
# >>> foo.pop = 2
# >>> ...
# >>> def test_foo():
# >>> assert foo().pop() == 2
#
# if we stop resolution when we see an LPAREN, we resolve `foo`
# successfully, fail on `pop` and everything is OK, but if we try to
# traverse the LPAREN, we get `foo.pop = 2` in our values, which is
# wrong
if toktype == tokenize.OP:
if tok not in (".", ")"):
self.doing_resolution = None
return
# CASE 2: "assert" statement
# assert statement was reached, reset
# return the start line (start = (startrow, startcol))
if tok == "assert":
self.seen_tokens.clear()
self.doing_resolution = None
return start[0]
# handle tokens
# CASE 3: a name is being resolved,
# there is a previous token,
# and it's a "." operator
if (
self.doing_resolution
and prior_tok
and (prior_tok[0] == tokenize.OP and prior_tok[1] == ".")
):
# unpack and look for the attribute
obj, name = self.doing_resolution
if hasattr(obj, tok):
obj = getattr(obj, tok)
name = name + "." + tok
self.doing_resolution = (obj, name)
self.seen_tokens[name] = obj
# if we couldn't find a relevant attribute, reset on resolution so
# that we can try afresh
else:
self.doing_resolution = None
# CASE 4: a name is being resolved and there is no preceding "." or
# resolution was explicitly stopped
else:
# skip tokens we've seen, but grab them as the current things under
# resolution
if tok in self.seen_tokens:
self.doing_resolution = (self.seen_tokens[tok], tok)
return
# we've never seen this token before
else:
# try to resolve to a value
try:
value = self.frame_locals[tok]
except KeyError:
try:
value = self.frame_globals[tok]
except KeyError:
# unresolveable name -- short circuit
# shows up in some cases like `f().x` in which `x`
# might not be a name bound to a value
return
# add it (so we don't try it again unless we hit a new assert
# and reset)
self.seen_tokens[tok] = value
self.doing_resolution = (value, tok)