# Source code for nose2.plugins.prettyassert

    """
    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
            """
            # if the trace is a string, that means it was pickled/unpickled via the mp
            # plugin, and is not a real traceback -- we won't be able to do anything with it
            if isinstance(trace, str):
                return

            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)