mypyc/test/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 Callable, Iterator

from mypy import build
from mypy.errors import CompileError
from mypy.main import process_options
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)

    flags = re.search("# flags: (.*)$", program_text, flags=re.MULTILINE)

    # 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, 10))
    if flags:
        flag_list = flags.group(1).split()
        _, options = process_options(flag_list, require_targets=False)
    else:
        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, 10)
    options.export_types = True
    options.preserve_asts = True
    options.allow_empty_bodies = True
    options.strict_bytes = True
    options.disable_bytearray_promotion = True
    options.disable_memoryview_promotion = True
    options.per_module_options["__main__"] = {"mypyc": True}

    source = build.BuildSource("main", "__main__", program_text)
    # Construct input as a single source.
    # Parse and type check the input program.
    result = build.build(sources=[source], options=options, alt_lib_path=test_temp_dir)
    result.manager.metastore.close()
    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:
    # Skip linker outputs. Pushing them back combines with write_file's
    # +1 sec bump on .c files to make .c always newer than .so, forcing
    # an unconditional rebuild that would mask Extension.depends bugs.
    # See setuptools/_distutils/command/build_ext.py:`build_extension`.
    for dirpath, _, filenames in os.walk(dir):
        for name in filenames:
            if name.endswith((".so", ".pyd", ".o", ".obj")):
                continue
            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_10* (or for any Python version):
          Use Python 3.10+ 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, 10), strict_traceback_checks=True
    )
    # A suffix like _python3_10 is used to set the target C API version.
    m = re.search(r"_python([0-9]+)_([0-9]+)(_|\b)", name)
    if m:
        version = (int(m.group(1)), int(m.group(2)))
        assert version >= (3, 10), f"Unsupported _python* suffix: {name}"
        options.capi_version = version
        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}"
    if has_test_name_tag(name, "experimental"):
        options.experimental_features = True
    return options


def has_test_name_tag(name: str, tag: str) -> bool:
    """Check if a test case name contains a tag token like ``_experimental``.

    A tag matches if it appears as a full underscore-delimited token:
    ``foo_tag_bar`` or ``foo_tag``.
    """
    return re.search(rf"(?:^|_){re.escape(tag)}(?:_|$)", name) is not None