File: _llvm.py

package info (click to toggle)
python3.14 3.14.2-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 128,000 kB
  • sloc: python: 752,614; ansic: 717,151; xml: 31,250; sh: 5,989; cpp: 4,063; makefile: 1,996; objc: 787; lisp: 502; javascript: 136; asm: 75; csh: 12
file content (121 lines) | stat: -rw-r--r-- 3,572 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
"""Utilities for invoking LLVM tools."""

import asyncio
import functools
import os
import re
import shlex
import subprocess
import typing

import _targets


_LLVM_VERSION = "21"
_EXTERNALS_LLVM_TAG = "llvm-19.1.7.0"

_P = typing.ParamSpec("_P")
_R = typing.TypeVar("_R")
_C = typing.Callable[_P, typing.Awaitable[_R]]


def _async_cache(f: _C[_P, _R]) -> _C[_P, _R]:
    cache = {}
    lock = asyncio.Lock()

    @functools.wraps(f)
    async def wrapper(
        *args: _P.args, **kwargs: _P.kwargs  # pylint: disable = no-member
    ) -> _R:
        async with lock:
            if args not in cache:
                cache[args] = await f(*args, **kwargs)
            return cache[args]

    return wrapper


_CORES = asyncio.BoundedSemaphore(os.cpu_count() or 1)


async def _run(tool: str, args: typing.Iterable[str], echo: bool = False) -> str | None:
    command = [tool, *args]
    async with _CORES:
        if echo:
            print(shlex.join(command))
        try:
            process = await asyncio.create_subprocess_exec(
                *command, stdout=subprocess.PIPE
            )
        except FileNotFoundError:
            return None
        out, _ = await process.communicate()
    if process.returncode:
        raise RuntimeError(f"{tool} exited with return code {process.returncode}")
    return out.decode()


@_async_cache
async def _check_tool_version(
    name: str, llvm_version: str, *, echo: bool = False
) -> bool:
    output = await _run(name, ["--version"], echo=echo)
    _llvm_version_pattern = re.compile(rf"version\s+{llvm_version}\.\d+\.\d+\S*\s+")
    return bool(output and _llvm_version_pattern.search(output))


@_async_cache
async def _get_brew_llvm_prefix(llvm_version: str, *, echo: bool = False) -> str | None:
    output = await _run("brew", ["--prefix", f"llvm@{llvm_version}"], echo=echo)
    return output and output.removesuffix("\n")


@_async_cache
async def _find_tool(tool: str, llvm_version: str, *, echo: bool = False) -> str | None:
    # Unversioned executables:
    path = tool
    if await _check_tool_version(path, llvm_version, echo=echo):
        return path
    # Versioned executables:
    path = f"{tool}-{llvm_version}"
    if await _check_tool_version(path, llvm_version, echo=echo):
        return path
    # PCbuild externals:
    externals = os.environ.get("EXTERNALS_DIR", _targets.EXTERNALS)
    path = os.path.join(externals, _EXTERNALS_LLVM_TAG, "bin", tool)
    if await _check_tool_version(path, llvm_version, echo=echo):
        return path
    # Homebrew-installed executables:
    prefix = await _get_brew_llvm_prefix(llvm_version, echo=echo)
    if prefix is not None:
        path = os.path.join(prefix, "bin", tool)
        if await _check_tool_version(path, llvm_version, echo=echo):
            return path
    # Nothing found:
    return None


async def maybe_run(
    tool: str,
    args: typing.Iterable[str],
    echo: bool = False,
    llvm_version: str = _LLVM_VERSION,
) -> str | None:
    """Run an LLVM tool if it can be found. Otherwise, return None."""

    path = await _find_tool(tool, llvm_version, echo=echo)
    return path and await _run(path, args, echo=echo)


async def run(
    tool: str,
    args: typing.Iterable[str],
    echo: bool = False,
    llvm_version: str = _LLVM_VERSION,
) -> str:
    """Run an LLVM tool if it can be found. Otherwise, raise RuntimeError."""

    output = await maybe_run(tool, args, echo=echo, llvm_version=llvm_version)
    if output is None:
        raise RuntimeError(f"Can't find {tool}-{llvm_version}!")
    return output