From 384f7ff5aa5cb6a12d650a513975f7bff63df30b Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Mon, 27 Jul 2026 23:32:06 +0530 Subject: [PATCH] Make function more easily testable --- kitty/shaders/slang.py | 198 +++++++++++++++++++++-------------------- 1 file changed, 103 insertions(+), 95 deletions(-) diff --git a/kitty/shaders/slang.py b/kitty/shaders/slang.py index 6cf9b65de..49ad0aadc 100644 --- a/kitty/shaders/slang.py +++ b/kitty/shaders/slang.py @@ -891,107 +891,115 @@ def custom_shader(name: str = '') -> tuple[str, bytes, bytes]: return name, src, key(src) +def _build_custom_shader_pipeline_ir(slot: str, shaders: Iterable[str], cache_dir: str) -> str: + slot_module_name = f'{slot.replace("-", "_")}' + cache_dir = os.path.join(cache_dir, 'c') + ensure_cache_dir(cache_dir) + slot_dir = os.path.join(cache_dir, 'slots') + libdir = os.path.join(cache_dir, 'lib') + bc = list(slangc()) + ['-warnings-as-errors', 'all', '-lang', 'slang', '-I', libdir] + _, ct_shader, ct_key = custom_shader() + os.makedirs(libdir, exist_ok=True) + os.makedirs(slot_dir, exist_ok=True) + j = partial(os.path.join, libdir) + cache_ok = False + mtime = 0 + with suppress(FileNotFoundError), open(j('ct.key'), 'rb') as f: + cache_ok = f.read() == ct_key + mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) + if not cache_ok: + cp = subprocess.run( + bc + ['-module-name', 'kitty_custom_shader_types', '-o', j('kitty-custom-shader-types.slang-module'), '--', '-'], + input=ct_shader, + capture_output=True, + ) + if cp.returncode != 0: + raise SlangFailed('custom-types.slang', cp) + with open(j('ct.key'), 'wb') as f: + f.write(ct_key) + mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) + shaders = tuple(shaders) + module_names = {} + shaders_content_key = b'' + for name in shaders: + path, src, content_key = custom_shader(name) + shaders_content_key += b':' + content_key + path_key = key(path) + path_key_file = j(path_key.decode()) + '.key' + cache_ok = False + module_names[name] = modname = 'm' + content_key.decode() + with suppress(FileNotFoundError), open(path_key_file, 'rb') as f: + cache_ok = f.read() == content_key + mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) + if not cache_ok: + inc = ['-I', os.path.dirname(path)] if os.sep in path else [] + cp = subprocess.run( + bc + inc + ['-module-name', modname, '-o', j(f'{modname}.slang-module'), '--', '-'], + input=src, + capture_output=True, + ) + if cp.returncode != 0: + raise SlangFailed(name, cp) + with open(path_key_file, 'wb') as f: + f.write(content_key) + mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) + shaders_content_key += b':' + str(mtime).encode() + j = partial(os.path.join, slot_dir) + cache_ok = False + wrappers = {} + entry_points = [] + for i, name in enumerate(shaders): + module_name = module_names[name] + entry_point = f'fragment_main{i}' + wrapper_src = textwrap.dedent(f""" + #language slang 2026 + implementing {slot_module_name}; + import kitty_custom_shader_types; + import {module_name}; + + public float4 {entry_point}(float4 inp, KittyCustomShaderData d) {{ return fragment_main(inp, d); }} + """) + wrappers[f'wrapper{i}.slang'] = wrapper_src + entry_points.append(entry_point) + mod_src = get_custom_shader_src('pipeline').decode() + mod_src = mod_src.replace('// IMPORTS', '\n'.join(f'__include "{w}";' for w in wrappers), 1) + mod_src = mod_src.replace('// PIPELINE', '\n'.join(f'color = {w}(color, d);' for w in entry_points), 1) + # subprocess.run(['bat', '-P', '-l', 'cpp'], input=mod_src.encode()) + slot_key = key(slot, mod_src, shaders_content_key) + + with suppress(FileNotFoundError), open(j(f'{slot}.key'), 'rb') as f: + cache_ok = f.read() == slot_key + ans = os.path.join(slot_dir, f'{slot}.slang-module') + if cache_ok: + return ans + with tempfile.TemporaryDirectory() as tdir: + for wrapper_name, wrapper_src in wrappers.items(): + with open(os.path.join(tdir, wrapper_name), 'w') as f: + f.write(wrapper_src) + cp = subprocess.run( + bc + ['-I', tdir, '-module-name', slot_module_name, '-o', ans, '--', '-'], + cwd=tdir, + capture_output=True, + input=mod_src.encode(), + ) + if cp.returncode != 0: + raise SlangFailed(f'{slot}.slang', cp) + + with open(j(f'{slot}.key'), 'wb') as f: + f.write(slot_key) + return ans + + def build_custom_shader_pipeline_ir(slot: str = 'after-window-background', shaders: Iterable[str] = ('sample',), cache_dir: str = '') -> str: import kitty.constants as kc cache_dir = os.path.join(cache_dir or kc.cache_dir(), 'shaders') os.makedirs(cache_dir, exist_ok=True) - slot_module_name = f'{slot.replace("-", "_")}' - with lock_with_file(os.path.join(cache_dir, 'lock')): - cache_dir = os.path.join(cache_dir, 'c') - ensure_cache_dir(cache_dir) - slot_dir = os.path.join(cache_dir, 'slots') - libdir = os.path.join(cache_dir, 'lib') - bc = list(slangc()) + ['-warnings-as-errors', 'all', '-lang', 'slang', '-I', libdir] - _, ct_shader, ct_key = custom_shader() - os.makedirs(libdir, exist_ok=True) - os.makedirs(slot_dir, exist_ok=True) - j = partial(os.path.join, libdir) - cache_ok = False - mtime = 0 - with suppress(FileNotFoundError), open(j('ct.key'), 'rb') as f: - cache_ok = f.read() == ct_key - mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) - if not cache_ok: - cp = subprocess.run( - bc + ['-module-name', 'kitty_custom_shader_types', '-o', j('kitty-custom-shader-types.slang-module'), '--', '-'], - input=ct_shader, - capture_output=True, - ) - if cp.returncode != 0: - raise SlangFailed('custom-types.slang', cp) - with open(j('ct.key'), 'wb') as f: - f.write(ct_key) - mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) - shaders = tuple(shaders) - module_names = {} - shaders_content_key = b'' - for name in shaders: - path, src, content_key = custom_shader(name) - shaders_content_key += b':' + content_key - path_key = key(path) - path_key_file = j(path_key.decode()) + '.key' - cache_ok = False - module_names[name] = modname = 'm' + content_key.decode() - with suppress(FileNotFoundError), open(path_key_file, 'rb') as f: - cache_ok = f.read() == content_key - mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) - if not cache_ok: - inc = ['-I', os.path.dirname(path)] if os.sep in path else [] - cp = subprocess.run( - bc + inc + ['-module-name', modname, '-o', j(f'{modname}.slang-module'), '--', '-'], - input=src, - capture_output=True, - ) - if cp.returncode != 0: - raise SlangFailed(name, cp) - with open(path_key_file, 'wb') as f: - f.write(content_key) - mtime = max(mtime, os.fstat(f.fileno()).st_mtime_ns) - shaders_content_key += b':' + str(mtime).encode() - j = partial(os.path.join, slot_dir) - cache_ok = False - wrappers = {} - entry_points = [] - for i, name in enumerate(shaders): - module_name = module_names[name] - entry_point = f'fragment_main{i}' - wrapper_src = textwrap.dedent(f""" - #language slang 2026 - implementing {slot_module_name}; - import kitty_custom_shader_types; - import {module_name}; - - public float4 {entry_point}(float4 inp, KittyCustomShaderData d) {{ return fragment_main(inp, d); }} - """) - wrappers[f'wrapper{i}.slang'] = wrapper_src - entry_points.append(entry_point) - mod_src = get_custom_shader_src('pipeline').decode() - mod_src = mod_src.replace('// IMPORTS', '\n'.join(f'__include "{w}";' for w in wrappers), 1) - mod_src = mod_src.replace('// PIPELINE', '\n'.join(f'color = {w}(color, d);' for w in entry_points), 1) - # subprocess.run(['bat', '-P', '-l', 'cpp'], input=mod_src.encode()) - slot_key = key(slot, mod_src, shaders_content_key) - - with suppress(FileNotFoundError), open(j(f'{slot}.key'), 'rb') as f: - cache_ok = f.read() == slot_key - if not cache_ok: - with tempfile.TemporaryDirectory() as tdir: - for wrapper_name, wrapper_src in wrappers.items(): - with open(os.path.join(tdir, wrapper_name), 'w') as f: - f.write(wrapper_src) - cp = subprocess.run( - bc + ['-I', tdir, '-module-name', slot_module_name, '-o', j(f'{slot}.slang-module'), '--', '-'], - cwd=tdir, - capture_output=True, - input=mod_src.encode(), - ) - if cp.returncode != 0: - raise SlangFailed(f'{slot}.slang', cp) - - with open(j(f'{slot}.key'), 'wb') as f: - f.write(slot_key) - return os.path.join(slot_dir, f'{slot}.slang-module') + with lock_with_file( + os.path.join(cache_dir, 'lock'), + ): + return _build_custom_shader_pipeline_ir(slot, shaders, cache_dir) def test_slang_build() -> None: