prettyassert.py

"""
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.

"""

from __future__ import annotations

import collections
import inspect
import io
import re
import textwrap
import tokenize
import typing as t

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) -> None:
        # 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: dict[str, t.Any] = {}

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