assert python example file by print() and comments

assert_example.py

import importlib
import inspect
import os
import re
import sys
from argparse import ArgumentParser, Namespace
from difflib import Differ
from pathlib import Path
from typing import Callable, List
from unittest.mock import patch

STRIP_PATTERN: str = r'^[#> ]+'
SINGLE_LINE_COMMENT: str = r'#+'
MULTILINE_COMMENT: str = r"'{3}|^\"{3}"


class InvalidComment(Exception):
    def __init__(self, file: Path, actual: List[str], expected: List[str]) -> None:
        self.file = file
        self._actual: List[str] = actual
        self._expected: List[str] = expected

    def diff(self) -> str:
        d = Differ()
        return '\n'.join(d.compare(self._actual, self._expected))

    @property
    def actual(self) -> str:
        return '\n'.join(self._actual)

    @property
    def expected(self) -> str:
        return '\n'.join(self._expected)


def create_mock_print(target_file: Path) -> Callable:
    def mock_print(*args, **kwargs):
        frame = inspect.currentframe().f_back.f_back.f_back
        if not target_file.samefile(frame.f_code.co_filename):
            return
        expected_comment: List[str] = []
        single_line: bool = False
        multi_line: bool = False
        target_line = frame.f_lineno
        for line in target_file.read_text().splitlines()[target_line:]:
            if re.match(SINGLE_LINE_COMMENT, line):
                if not single_line:
                    single_line = True
                expected_comment.append(re.sub(STRIP_PATTERN, '', line))
            elif single_line:
                break
            elif re.match(MULTILINE_COMMENT, line):
                if multi_line:
                    break
                multi_line = True
            else:
                if multi_line:
                    expected_comment.append(re.sub(STRIP_PATTERN, '', line))
                else:
                    break
        if not single_line and not multi_line:
            raise Exception(
                f'NotFound expected comment. target_line is {target_line + 1}'
            )
        actual = str(args[0]).splitlines()
        try:
            assert actual == expected_comment
        except AssertionError:
            raise InvalidComment(target_file, actual, expected_comment)

    return mock_print


def assert_example(file: Path):
    module_name = '.'.join(str(file.with_suffix('')).split('/'))
    with patch('builtins.print') as mock_print:
        mock_print.side_effect = create_mock_print(file)
        importlib.import_module(module_name)


def main():
    arg_parser = ArgumentParser()
    arg_parser.add_argument(
        '--input', help='input python file or a directory', type=str
    )
    namespace: Namespace = arg_parser.parse_args()
    sys.path.append(os.getcwd())
    target_path = Path(namespace.input)

    try:
        if target_path.is_file():
            assert_example(target_path)
        elif target_path.is_dir():
            for file in target_path.glob('**/*.py'):
                assert_example(file)
                sys.stdout.write('.')
        else:
            raise Exception('Unexpected path type. accept only file or directory')
    except InvalidComment as e:
        sys.stderr.write(f'Diff: \n')
        sys.stderr.write(f'{e.diff()}\n\n')
        sys.stderr.write(f'Actual: \n')
        sys.stderr.write(f'{e.actual}\n\n')
        sys.stderr.write(f'Failed: {str(e.file)}\n')
        exit(1)
    print('\nSuccess')


if __name__ == '__main__':
    main()
"""
The script asserts python example files by print() and comments.

~~~ invalid_example.py ~~~
print(f'foo bar {1+2}')
# foo samp 4
~~~ invalid_example.py ~~~

$ python3 assert_example.py --input example.py
Diff: 
- foo bar 3
+ foo samp 4

Actual: 
foo bar 3

Failed: invalid_example.py
"""