<reponame>i2mint/py2mint
"""Meta-interfaces """
import doctest

MAX_LINE_LENGTH = 72  # https://en.wikipedia.org/wiki/Characters_per_line


def _prefix_lines(s: str, prefix: str = '# ', even_if_empty: bool = False) -> str:
    r"""
    Prefix every line of s with given prefix.

    :param s: String whose lines you want to prefix.
    :param prefix: Desired prefix string. Default is '# ', to have the effect of "commenting out" line
    :param even_if_empty: Whether to prefix empty strings or not.
    :return: A string whose lines have been prefixed.

    >>> _prefix_lines('something to comment out')
    '# something to comment out'
    >>> _prefix_lines('A line you want to prefix', prefix='PREFIX: ')
    'PREFIX: A line you want to prefix'
    >>> print(_prefix_lines('What happens\nif the thing to comment out\nhas multiple lines?'))
    # What happens
    # if the thing to comment out
    # has multiple lines?
    >>> _prefix_lines('')  # see that an empty string is returned as is
    ''
    >>> _prefix_lines('', even_if_empty=True)  # unless you ask for it
    '# '
    """
    if not even_if_empty:
        if len(s) == 0:
            return s
    return '\n'.join(map(lambda x: prefix + x, s.split('\n')))


import doctest
from typing import Callable
import re
from inspect import getdoc

comment_strip_p = re.compile(r'(?m)^ *#.*\n?')

doctest_line_p = re.compile('\s*>>>')
empty_line = re.compile('\s*$')


def non_doctest_lines(doc):
    r"""Generator of lines of the doc string that are not in a doctest scope.

    >>> def _test_func():
    ...     '''Line 1
    ...     Another
    ...     >>> doctest_1
    ...     >>> doctest_2
    ...     line_after_a_doc_test
    ...     another_line_that_is_in_the_doc_test scope
    ...
    ...     But now we're out of a doctest's scope
    ...
    ...     >>> Oh no, another doctest!
    ...     '''
    >>> from inspect import getdoc
    >>>
    >>> list(non_doctest_lines(getdoc(_test_func)))
    ['Line 1', 'Another', "But now we're out of a doctest's scope", '']

    :param doc:
    :return:
    """
    last_line_was_a_doc_test = False
    for line in doc.splitlines():
        if not doctest_line_p.match(line):
            if not last_line_was_a_doc_test:
                yield line
                last_line_was_a_doc_test = False
            else:
                if empty_line.match(line):
                    last_line_was_a_doc_test = False
        else:
            last_line_was_a_doc_test = True


def strip_comments(code):
    code = str(code)
    return comment_strip_p.sub('', code)


def mk_example_wants_callback(source_want_func: Callable[[str, str], Callable]):
    def example_wants_callback(example, *args, **kwargs):
        want = example.want.strip()
        if want:
            source = example.source.strip()
            return source_want_func(source, want, *args, **kwargs)
        else:
            return example.source

    return example_wants_callback


def split_line_comments(s):
    t = s.split('#')
    if len(t) == 1:
        comment = ''
    else:
        s, comment = t
    return s, comment


def _assert_wants(source, want, wrap_func_name=None):
    is_a_multiline = len(source.split('\n')) > 1

    if not is_a_multiline:
        source, comment = split_line_comments(source)
        if wrap_func_name is None:
            t = f'({source}) == {want} #{comment}'
        else:
            t = f'{wrap_func_name}({source}) == {wrap_func_name}({want}) #{comment}'
        if "'" in t and not '"' in t:
            strchr = '"'
            return 'assert {t}, {strchr}{t}{strchr}'.format(t=t, strchr=strchr)
        elif '"' in t and not "'" in t:
            strchr = "'"
            return 'assert {t}, {strchr}{t}{strchr}'.format(t=t, strchr=strchr)
        else:
            return 'assert {t}'.format(t=t)
    else:  # if you didn't return before
        if wrap_func_name is None:
            return f'actual = {source}\nexpected = {want}\nassert actual == expected'
        else:
            return (
                f'actual = {wrap_func_name}({source})\nexpected = {wrap_func_name}({want})\n'
                'assert actual == expected'
            )


def _output_prefix(source, want, prefix='# OUTPUT: '):
    return source + '\n' + prefix + want + '\n'


output_prefix = mk_example_wants_callback(_output_prefix)
assert_wants = mk_example_wants_callback(_assert_wants)


def doctest_string_trans_lines(
    doctest_obj: doctest.DocTest, example_callback=assert_wants
):
    for example in doctest_obj.examples:
        yield example_callback(example)


def _doctest_string_gen(obj, example_callback, recurse=True):
    doctest_finder = doctest.DocTestFinder(verbose=False, recurse=recurse)
    doctest_objs = doctest_finder.find(obj)
    for doctest_obj in doctest_objs:
        yield from doctest_string_trans_lines(doctest_obj, example_callback)


def doctest_string(obj, example_callback=assert_wants, recurse=True):
    """
    Extract the doctests found in given object.
    :param obj: Object (module, class, function, etc.) you want to extract doctests from.
    :params output_prefix:
    :param recurse: Whether the process should find doctests in the attributes of the object, recursively.
    :return: A string containing the doctests, with output lines prefixed by '# Output:'
    """
    return '\n'.join(_doctest_string_gen(obj, example_callback, recurse=recurse))


from functools import partial

doctest_string.for_output_prefix = partial(
    doctest_string, example_callback=output_prefix
)
doctest_string.for_assert_wants = partial(doctest_string, example_callback=assert_wants)


def doctest_string_print(obj, example_callback=assert_wants, recurse=True):
    """
    Extract the doctests found in given object.
    :param obj: Object (module, class, function, etc.) you want to extract doctests from.
    :param recurse: Whether the process should find doctests in the attributes of the object, recursively.
    :return: A string containing the doctests, with output lines prefixed by '# Output:'
    """
    return print(doctest_string(obj, example_callback, recurse=recurse))


def old_doctest_string(
    obj, output_prefix='# OUTPUT: ', include_attr_without_doctests=False, recurse=True,
):
    """
    Extract the doctests found in given object.
    :param obj: Object (module, class, function, etc.) you want to extract doctests from.
    :param output_prefix:
    :param recurse: Whether the process should find doctests in the attributes of the object, recursively.
    :return: A string containing the doctests, with output lines prefixed by '# Output:'
    """
    doctest_finder = doctest.DocTestFinder(verbose=False, recurse=recurse)
    r = doctest_finder.find(obj)
    s = ''
    for rr in r:
        header = f'# {rr.name} '
        header += '#' * max(0, MAX_LINE_LENGTH - len(header)) + '\n'
        ss = ''
        for example in rr.examples:
            want = example.want
            want = want.strip()
            ss += '\n' + example.source + _prefix_lines(want, prefix=output_prefix)
        if include_attr_without_doctests:
            s += header + ss
        elif len(ss) > 0:  # only append this attr if ss is non-empty
            s += header + ss
    return s


# import sphinx

if __name__ == '__main__':
    print(doctest_string(_prefix_lines))
# # _prefix_lines ########################################################
#
# _prefix_lines('something to comment out')
# # OUTPUT: '# something to comment out'
# _prefix_lines('A line you want to prefix', prefix='PREFIX: ')
# # OUTPUT: 'PREFIX: A line you want to prefix'
# print(_prefix_lines('What happens\nif the thing to comment out\nhas multiple lines?'))
# # OUTPUT: # What happens
# # OUTPUT: # if the thing to comment out
# # OUTPUT: # has multiple lines?
# _prefix_lines('')  # see that an empty string is returned as is
# # OUTPUT: ''
# _prefix_lines('', even_if_empty=True)  # unless you ask for it
# # OUTPUT: '# '
