#!/usr/bin/env sage-python

"""
Given the output of doctest and a file, adjust the doctests so they won't fail.

Doctest failures due to exceptions are ignored.

AUTHORS::

- Nicolas M. Thiéry <nthiery at users dot sf dot net>  Initial version (2008?)

- Andrew Mathas <andrew dot mathas at sydney dot edu dot au> 2013-02-14
  Cleaned up the code and hacked it so that the script can now cope with the
  situations when either the expected output or computed output are empty.
  Added doctest to sage.tests.cmdline
"""
import os
import subprocess
from argparse import ArgumentParser, FileType
from sage.misc.temporary_file import tmp_filename

parser = ArgumentParser(description="Given an input file with doctests, this creates a modified file that passes the doctests (modulo any raised exceptions). By default, the input file is modified. You can also name an output file.")
parser.add_argument('-l', '--long',
                    dest='long', action="store_true", default=False)
parser.add_argument("input", help="input filename",
                    type=FileType('r'))
parser.add_argument('output', nargs='?', help="output filename",
                    type=FileType('w'))

args = parser.parse_args()

# set input and output files
test_file = args.input
src_in = test_file.read()
test_file.close()

# put the output of the test into sage's temporary directory
doc_file = tmp_filename()
os.system('sage -t %s %s > %s' % ('--long' if args.long else '',
                                  test_file.name, doc_file))
with open(doc_file, 'r') as doc:
    doc_out = doc.read()
sep = "**********************************************************************\n"

doctests = doc_out.split(sep)
src_in_lines = src_in.splitlines()

for block in doctests:
    if 'Failed example:' not in block:
        continue  # sanity checking, but shouldn't happen

    # Extract the line, what was expected, and was got.
    line = block.find('line ')             # block should contain:  'line ##, in ...', where ## is an integer
    comma = block.find(', in ')            # we try to extract the line number which gives the test failure
    if line == -1 or comma == -1:
        continue   # but if something goes wrong we give up
    line_num = eval(block[line + 5:comma])

    # Take care of multiline examples.
    if 'Expected' in block:
        i1 = block.index('Failed example')
        i2 = block.index('Expected')
        example_len = block[i1:i2].count('\n') - 1
        line_num += example_len - 1

    if 'Expected nothing' in block:
        exp = block.find('Expected nothing')
        block = '\n' + block[exp + 17:]  # so that split('\nGot:\n') does not fail below
    elif 'Expected:' in block:
        exp = block.find('Expected:\n')
        block = block[exp + 10:]
    elif 'Exception raised' in block:
        exp = block.find('Exception raised')
        block = '\nGot:\n' + block[exp + 17:]
    else:
        continue
    # Error testing.
    if 'Traceback (most recent call last):' in block:
        expected, got = block.split('\nGot:\n')
        got = got.splitlines()
        got = ['    Traceback (most recent call last):', '    ...', got[-1]]
    elif block[-21:] == 'Got:\n    <BLANKLINE>\n':
        expected = block[:-22]
        got = ['']
    else:
        expected, got = block.split('\nGot:\n')
        got = got.splitlines()      # got can't be the empty string

    # If we expected nothing, and got something, then we need to insert the line before line_num
    # and match indentation with line number line_num-1
    if expected == '':
        indent = (len(src_in_lines[line_num - 1]) - len(src_in_lines[line_num - 1].lstrip()))
        src_in_lines[line_num - 1] += '\n' + '\n'.join('%s%s' % (' ' * indent, line.lstrip()) for line in got)
        continue

    # Guess how much extra indenting got needs to match with the indentation
    # of src_in_lines - we match the indentation with the line in ``got`` which
    # has the smallest indentation after lstrip(). Note that the amount of indentation
    # required could be negative if the ``got`` block is indented. In this case
    # ``indent`` is set to zero.
    indent = max(0, (len(src_in_lines[line_num]) - len(src_in_lines[line_num].lstrip()) - min(len(got[j]) - len(got[j].lstrip()) for j in range(len(got)))))

    expected = expected.splitlines()

    # Double check that what was expected was indeed in the source file and if
    # it is not then then print a warning for the user which contains the
    # problematic lines.
    if any(expected[i].strip() != src_in_lines[line_num + i].strip()
           for i in range(len(expected))):
        import warnings
        txt = "Did not manage to replace\n%s\n%s\n%s\nwith\n%s\n%s\n%s"
        warnings.warn(txt % ('>' * 40, '\n'.join(expected), '>' * 40,
                             '<' * 40, '\n'.join(got), '<' * 40))
        continue

    # If we got something when we expected nothing then we delete the line from the
    # output, otherwise, add all of what we `got` onto the end of src_in_lines[line_num]
    if got == ['']:
        src_in_lines[line_num] = None
    else:
        src_in_lines[line_num] = '\n'.join((' ' * indent + got[i])
                                           for i in range(len(got)))

    # Mark any remaining `expected` lines as ``None`` so as to preserve the line numbering
    for i in range(1, len(expected)):
        src_in_lines[line_num + i] = None

# Overwrite the source (or output file specified on the command line)
if args.output:
    test_output = args.output
else:
    test_output = open(test_file.name, 'w')

for line in src_in_lines:
    if line is None:
        continue
    test_output.write(line)
    test_output.write('\n')

test_output.close()

# Show summary of changes
if args.output:
    print('The fixed doctests have been saved as {0}.'.format(test_output.name))
else:
    from sage.env import SAGE_ROOT
    relative = os.path.relpath(test_output.name, SAGE_ROOT)
    print(relative)
    if relative.startswith('..'):
        print('Fixed source file is not part of Sage.')
    else:
        subprocess.call(['git', 'diff', relative], cwd=SAGE_ROOT)
