example_import_handler.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "jsonnet==0.21.0rc2",
# ]
# ///
import _jsonnet
import os

# WARNING:
# This creates example input files in the location you run the script!
# It will leave the files there and you will have to clean them up manually!

SEARCH_PATHS = [
    "dirA",
    "dirB/one",
    "dirB/two",
]

EXAMPLE_INPUTS = {
    "dirA/main.jsonnet": """
{
   "a": import "a.jsonnet",
   "b": import "b.jsonnet",
   "c": import "c.jsonnet",
}
    """.strip(),
    "dirA/b.jsonnet": "'this is b.jsonnet in dirA'",
    "dirB/one/a.jsonnet": "'this is a.jsonnet in dirB/one'",
    "dirB/one/c.jsonnet": "{'sub': import 'd.jsonnet', 'text': 'this is c.jsonnet in dirB/one'}",
    "dirB/two/d.jsonnet": "'this is d.jsonnet in dirB/two'",
    "dirB/two/b.jsonnet": "'this is b.jsonnet in dirB/two - hidden by the other b.jsonnet'",
}

def load_import_from_path(base_dir, import_path):
    # Handle an import of `import_path`, found in a file in `base_dir`.
    print(f'import {import_path} (from {base_dir})')
    # Look for the import_path within each of the SEARCH_PATHS entries.
    for sp in SEARCH_PATHS:
        path = os.path.join(sp, import_path)
        try:
            with open(path, 'rb') as f:
                content = f.read()
                print(f'found: {path} ({len(content)} bytes)')
                # Return the path to the file we found, and its content.
                return path, content
        except FileNotFoundError as e:
            # Trace what we're doing
            print(f'not found: {path}: {e}')
    print(f'no matching file found for import {import_path} (from {base_dir})')

def setup_example_inputs():
    for sp in SEARCH_PATHS:
        os.makedirs(sp, exist_ok=True)
    for path, content in EXAMPLE_INPUTS.items():
        with open(path, 'wb') as f:
            f.write(content.encode('utf-8'))

def main():
    setup_example_inputs()
    print('Using custom import_callback:')
    res1 = _jsonnet.evaluate_file('dirA/main.jsonnet', import_callback=load_import_from_path)
    print(res1)
    print('Using jpathdir:')
    res2 = _jsonnet.evaluate_file('dirA/main.jsonnet', jpathdir=SEARCH_PATHS)
    print(res2)

if __name__ == "__main__":
    main()