Function extract_triton_code

export_hf.py:237–262

Extract the Triton kernel code from a Python file. Returns everything from the first import statement onward, skipping the module docstring and KERNEL_TYPE/BACKEND declarations.

(source: str)

# ---------------------------------------------------------------------------

def extract_triton_code(source: str) -> str:
    """
    Extract the Triton kernel code from a Python file.

    Returns everything from the first import statement onward, skipping
    the module docstring and KERNEL_TYPE/BACKEND declarations.
    """
    lines = source.split("\n")

    # Find the first import line
    import_idx = None
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            import_idx = i
            break

    if import_idx is not None:
        return "\n".join(lines[import_idx:])

    # Fallback: return everything after KERNEL_TYPE line
    for i, line in enumerate(lines):
        if line.strip().startswith("KERNEL_TYPE"):
            return "\n".join(lines[i + 1 :])

    return source


# ---------------------------------------------------------------------------