testutil.py

"""Helpers for writing tests"""

from __future__ import annotations

import contextlib
import os
import os.path
import re
import shutil
from collections.abc import Iterator
from typing import Callable

from mypy import build
from mypy.errors import CompileError
from mypy.nodes import Expression, MypyFile
from mypy.options import Options
from mypy.test.config import test_temp_dir
from mypy.test.data import DataDrivenTestCase, DataSuite
from mypy.test.helpers import assert_string_arrays_equal
from mypy.types import Type
from mypyc.analysis.ircheck import assert_func_ir_valid
from mypyc.common import IS_32_BIT_PLATFORM, PLATFORM_SIZE
from mypyc.errors import Errors
from mypyc.ir.func_ir import FuncIR
from mypyc.ir.module_ir import ModuleIR
from mypyc.irbuild.main import build_ir
from mypyc.irbuild.mapper import Mapper
from mypyc.options import CompilerOptions
from mypyc.test.config import test_data_prefix

# The builtins stub used during icode generation test cases.
ICODE_GEN_BUILTINS = os.path.join(test_data_prefix, "fixtures/ir.py")
# The testutil support library
TESTUTIL_PATH = os.path.join(test_data_prefix, "fixtures/testutil.py")


class MypycDataSuite(DataSuite):
    # Need to list no files, since this will be picked up as a suite of tests
    files: list[str] = []
    data_prefix = test_data_prefix


def builtins_wrapper(
    func: Callable[[DataDrivenTestCase], None], path: str
) -> Callable[[DataDrivenTestCase], None]:
    """Decorate a function that implements a data-driven test case to copy an
    alternative builtins module implementation in place before performing the
    test case. Clean up after executing the test case.
    """
    return lambda testcase: perform_test(func, path, testcase)


@contextlib.contextmanager
def use_custom_builtins(builtins_path: str, testcase: DataDrivenTestCase) -> Iterator[None]:
    for path, _ in testcase.files:
        if os.path.basename(path) == "builtins.pyi":
            default_builtins = False
            break
    else:
        # Use default builtins.
        builtins = os.path.abspath(os.path.join(test_temp_dir, "builtins.pyi"))
        shutil.copyfile(builtins_path, builtins)
        default_builtins = True

    # Actually perform the test case.
    try:
        yield None
    finally:
        if default_builtins:
            # Clean up.
            os.remove(builtins)


def perform_test(
    func: Callable[[DataDrivenTestCase], None], builtins_path: str, testcase: DataDrivenTestCase
) -> None:
    for path, _ in testcase.files:
        if os.path.basename(path) == "builtins.py":
            default_builtins = False
            break
    else:
        # Use default builtins.
        builtins = os.path.join(test_temp_dir, "builtins.py")
        shutil.copyfile(builtins_path, builtins)
        default_builtins = True

    # Actually perform the test case.
    func(testcase)

    if default_builtins:
        # Clean up.
        os.remove(builtins)


def build_ir_for_single_file(
    input_lines: list[str], compiler_options: CompilerOptions | None = None
) -> list[FuncIR]:
    return build_ir_for_single_file2(input_lines, compiler_options)[0].functions


def build_ir_for_single_file2(
    input_lines: list[str], compiler_options: CompilerOptions | None = None
) -> tuple[ModuleIR, MypyFile, dict[Expression, Type], Mapper]:
    program_text = "\n".join(input_lines)

    # By default generate IR compatible with the earliest supported Python C API.
    # If a test needs more recent API features, this should be overridden.
    compiler_options = compiler_options or CompilerOptions(capi_version=(3, 9))
    options = Options()
    options.show_traceback = True
    options.hide_error_codes = True
    options.use_builtins_fixtures = True
    options.strict_optional = True
    options.python_version = compiler_options.python_version or (3, 8)
    options.export_types = True
    options.preserve_asts = True
    options.allow_empty_bodies = True
    options.per_module_options["__main__"] = {"mypyc": True}

    source = build.BuildSource("main", "__main__", program_text)
    # Construct input as a single single.
    # Parse and type check the input program.
    result = build.build(sources=[source], options=options, alt_lib_path=test_temp_dir)
    if result.errors:
        raise CompileError(result.errors)

    errors = Errors(options)
    mapper = Mapper({"__main__": None})
    modules = build_ir(
        [result.files["__main__"]], result.graph, result.types, mapper, compiler_options, errors
    )
    if errors.num_errors:
        raise CompileError(errors.new_messages())

    module = list(modules.values())[0]
    for fn in module.functions:
        assert_func_ir_valid(fn)
    tree = result.graph[module.fullname].tree
    assert tree is not None
    return module, tree, result.types, mapper


def update_testcase_output(testcase: DataDrivenTestCase, output: list[str]) -> None:
    # TODO: backport this to mypy
    assert testcase.old_cwd is not None, "test was not properly set up"
    testcase_path = os.path.join(testcase.old_cwd, testcase.file)
    with open(testcase_path) as f:
        data_lines = f.read().splitlines()

    # We can't rely on the test line numbers to *find* the test, since
    # we might fix multiple tests in a run. So find it by the case
    # header. Give up if there are multiple tests with the same name.
    test_slug = f"[case {testcase.name}]"
    if data_lines.count(test_slug) != 1:
        return
    start_idx = data_lines.index(test_slug)
    stop_idx = start_idx + 11
    while stop_idx < len(data_lines) and not data_lines[stop_idx].startswith("[case "):
        stop_idx += 1

    test = data_lines[start_idx:stop_idx]
    out_start = test.index("[out]")
    test[out_start + 1 :] = output
    data_lines[start_idx:stop_idx] = test + [""]
    data = "\n".join(data_lines)

    with open(testcase_path, "w") as f:
        print(data, file=f)


def assert_test_output(
    testcase: DataDrivenTestCase,
    actual: list[str],
    message: str,
    expected: list[str] | None = None,
    formatted: list[str] | None = None,
) -> None:
    __tracebackhide__ = True

    expected_output = expected if expected is not None else testcase.output
    if expected_output != actual and testcase.config.getoption("--update-data", False):
        update_testcase_output(testcase, actual)

    assert_string_arrays_equal(
        expected_output, actual, f"{message} ({testcase.file}, line {testcase.line})"
    )


def get_func_names(expected: list[str]) -> list[str]:
    res = []
    for s in expected:
        m = re.match(r"def ([_a-zA-Z0-9.*$]+)\(", s)
        if m:
            res.append(m.group(1))
    return res


def remove_comment_lines(a: list[str]) -> list[str]:
    """Return a copy of array with comments removed.

    Lines starting with '--' (but not with '---') are removed.
    """
    r = []
    for s in a:
        if s.strip().startswith("--") and not s.strip().startswith("---"):
            pass
        else:
            r.append(s)
    return r


def print_with_line_numbers(s: str) -> None:
    lines = s.splitlines()
    for i, line in enumerate(lines):
        print("%-4d %s" % (i + 1, line))


def heading(text: str) -> None:
    print("=" * 20 + " " + text + " " + "=" * 20)


def show_c(cfiles: list[list[tuple[str, str]]]) -> None:
    heading("Generated C")
    for group in cfiles:
        for cfile, ctext in group:
            print(f"== {cfile} ==")
            print_with_line_numbers(ctext)
    heading("End C")


def fudge_dir_mtimes(dir: str, delta: int) -> None:
    for dirpath, _, filenames in os.walk(dir):
        for name in filenames:
            path = os.path.join(dirpath, name)
            new_mtime = os.stat(path).st_mtime + delta
            os.utime(path, times=(new_mtime, new_mtime))


def replace_word_size(text: list[str]) -> list[str]:
    """Replace WORDSIZE with platform specific word sizes"""
    result = []
    for line in text:
        index = line.find("WORD_SIZE")
        if index != -1:
            # get 'WORDSIZE*n' token
            word_size_token = line[index:].split()[0]
            n = int(word_size_token[10:])
            replace_str = str(PLATFORM_SIZE * n)
            result.append(line.replace(word_size_token, replace_str))
        else:
            result.append(line)
    return result


def infer_ir_build_options_from_test_name(name: str) -> CompilerOptions | None:
    """Look for magic substrings in test case name to set compiler options.

    Return None if the test case should be skipped (always pass).

    Supported naming conventions:

      *_64bit*:
          Run test case only on 64-bit platforms
      *_32bit*:
          Run test caseonly on 32-bit platforms
      *_python3_8* (or for any Python version):
          Use Python 3.8+ C API features (default: lowest supported version)
      *StripAssert*:
          Don't generate code for assert statements
    """
    # If this is specific to some bit width, always pass if platform doesn't match.
    if "_64bit" in name and IS_32_BIT_PLATFORM:
        return None
    if "_32bit" in name and not IS_32_BIT_PLATFORM:
        return None
    options = CompilerOptions(strip_asserts="StripAssert" in name, capi_version=(3, 9))
    # A suffix like _python3_9 is used to set the target C API version.
    m = re.search(r"_python([3-9]+)_([0-9]+)(_|\b)", name)
    if m:
        options.capi_version = (int(m.group(1)), int(m.group(2)))
        options.python_version = options.capi_version
    elif "_py" in name or "_Python" in name:
        assert False, f"Invalid _py* suffix (should be _pythonX_Y): {name}"
    return options