ely-eye/bench/build_code.py

"""Build code sets from CodeSearchNet.

A function's docstring is the query and the function itself is the answer, so
the gold is the dataset's own (func_path_in_repository, func_name) pair. Two
hazards are handled explicitly: the published function text still contains its
docstring, which would hand BM25 a free exact match, and 672 docstrings are
shared by more than one function, which makes their gold ambiguous. Whole
repository clones are a b-5 concern, not part of this set.
"""

from __future__ import annotations

import ast
import hashlib
import json
import re
import warnings
from pathlib import Path
from typing import Any

from pyarrow.parquet import ParquetFile

from bench.corpus_sources import source_for
from bench.sets_common import (
    DATA_ROOT,
    SETS,
    corpus_path,
    corpus_record,
    query_record,
    safe_id,
    safe_segment,
    sample_by_holder,
    write_jsonl,
)

DATASET = "code-search-net/code_search_net"
SOURCE = source_for(DATASET)
REVISION = str(SOURCE["revision"])
SNAPSHOT_AT = str(SOURCE["snapshot_at"])
PARQUET = "python/test-00000-of-00001.parquet"
SCENARIO = "code-search"
HELD_OUT_BUCKET = 15


def strip_docstring(code: str) -> tuple[str, bool]:
    """Remove a function's own docstring lines; report whether that succeeded."""
    try:
        with warnings.catch_warnings():
            # Corpus code is third-party; its own escape sequences are not our problem.
            warnings.simplefilter("ignore", SyntaxWarning)
            tree = ast.parse(code)
    except (SyntaxError, ValueError):
        return code, False
    body = tree.body
    if len(body) == 1 and isinstance(body[0], ast.FunctionDef | ast.AsyncFunctionDef):
        body = body[0].body
    if not body or not isinstance(body[0], ast.Expr):
        return code, False
    first = body[0]
    value = first.value
    if not isinstance(value, ast.Constant) or not isinstance(value.value, str):
        return code, False
    lines = code.splitlines(keepends=True)
    start, end = first.lineno - 1, first.end_lineno
    stripped = "".join(lines[:start] + lines[end:])
    if not stripped.strip():
        return code, False
    return stripped, True


def normalize_docstring(text: str) -> str:
    return re.sub(r"\s+", " ", text).strip().lower()


def repo_slug(repository: str) -> str:
    """CodeSearchNet names are 'owner/name'; one safe form is used as holder and path."""
    return safe_id(repository.replace("/", "-"), f"repository_name {repository!r}")


def function_path(repository: str, func_path: str, func_name: str, tag: str) -> str:
    """Where the function text lives under DATA_ROOT, mirroring the repo layout.

    The tag keeps two paths that sanitize alike, or one file holding two functions
    with the same name, from overwriting each other.
    """
    segments = [repo_slug(repository)]
    segments += [
        safe_segment(segment, f"func_path_in_repository {func_path!r}")
        for segment in func_path.split("/")
        if segment not in ("", ".")
    ]
    stem = segments[-1].removesuffix(".py") or "module"
    segments[-1] = f"{stem}-{tag}"
    return (Path("codesearchnet") / Path(*segments)).with_suffix(".py").as_posix()


def read_rows() -> list[dict[str, Any]]:
    table = ParquetFile(corpus_path(DATASET, PARQUET)).read()
    return table.select(
        ["repository_name", "func_path_in_repository", "func_name", "func_code_string",
         "func_documentation_string", "func_code_url"]
    ).to_pylist()


def build_functions() -> tuple[list, list[dict[str, Any]], dict]:
    rows = read_rows()
    corpus: list[dict[str, Any]] = []
    candidates: list[dict[str, Any]] = []
    skipped = {"no_docstring": 0, "empty_after_strip": 0, "docstring_left_in_text": 0}

    for index, row in enumerate(rows):
        repository = str(row["repository_name"])
        func_path = str(row["func_path_in_repository"])
        func_name = str(row["func_name"])
        code = str(row["func_code_string"])
        docstring = str(row["func_documentation_string"]).strip()
        if not docstring or not code.strip():
            skipped["no_docstring"] += 1
            continue
        stripped, removed = strip_docstring(code)
        if not removed and normalize_docstring(docstring) in normalize_docstring(stripped):
            skipped["docstring_left_in_text"] += 1
            continue
        if not stripped.strip():
            skipped["empty_after_strip"] += 1
            continue
        # One dataset row is one function; its index is the only thing that makes the
        # identity unique, because a file may define the same name twice.
        tag = hashlib.sha256(f"{func_path}#{func_name}#{index}".encode()).hexdigest()[:10]
        relative = function_path(repository, func_path, func_name, tag)
        target = DATA_ROOT / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(stripped, encoding="utf-8")
        source_id = f"csn:{repository}:{func_path}:{func_name}#{index}"
        corpus.append(
            corpus_record(
                source_id=source_id,
                modality="code",
                path=relative,
                title=f"{func_path}:{func_name}",
                observed_at=SNAPSHOT_AT,
                holder=repo_slug(repository),
                scenario=SCENARIO,
                root="derived",
            )
        )
        candidates.append(
            {
                "source_id": source_id,
                "repository": repo_slug(repository),
                "handle": f"{func_name}-{tag}",
                "url": str(row["func_code_url"]),
                "docstring": docstring,
                "key": normalize_docstring(docstring),
                "region": {"path": func_path, "symbol": func_name},
                "held_out": int(hashlib.sha256(source_id.encode()).hexdigest(), 16) % 100
                < HELD_OUT_BUCKET,
            }
        )

    keys: dict[str, int] = {}
    for candidate in candidates:
        keys[candidate["key"]] = keys.get(candidate["key"], 0) + 1
    ambiguous = sum(1 for candidate in candidates if keys[candidate["key"]] > 1)
    return corpus, candidates, {
        "functions_in_split": len(rows),
        "functions_kept": len(candidates),
        "skipped": skipped,
        "ambiguous_docstrings_dropped": ambiguous,
        "held_out_share": HELD_OUT_BUCKET,
    }


def to_query(candidate: dict[str, Any], unanswerable: bool) -> dict[str, Any]:
    """Turn one CodeSearchNet function into a docstring-to-code query."""
    suffix = "" if not unanswerable else "@held-out"
    origin = {
        "dataset": DATASET,
        "revision": REVISION,
        "item": f"url:{candidate['url']}",
    }
    if unanswerable:
        origin["construction"] = "held-out function"
    return query_record(
        query_id=f"csn:{candidate['repository']}:{candidate['handle']}{suffix}",
        modality="code",
        question=candidate["docstring"],
        answer=None if unanswerable else candidate["source_id"],
        answer_format="code",
        gold=[] if unanswerable else [{"source_id": candidate["source_id"], "region": candidate["region"]}],
        holder=candidate["repository"],
        scenario=SCENARIO,
        as_of=SNAPSHOT_AT,
        origin=origin,
    )


def build() -> tuple[list, list, dict]:
    corpus, candidates, stats = build_functions()
    answerable = sample_by_holder(
        [to_query(c, False) for c in candidates if not c["held_out"]], per_holder=2
    )
    unanswerable = [to_query(c, True) for c in candidates if c["held_out"]]
    queries = answerable + unanswerable[: len(answerable)]

    report = stats | {
        "datasets": [DATASET],
        "revision": REVISION,
        "holders": sorted({record["holder"] for record in corpus}),
        "queries": len(queries),
        "queries_answerable": len(answerable),
        "queries_held_out_unanswerable": len(queries) - len(answerable),
    }
    return corpus, queries, report


def main() -> None:
    corpus, queries, report = build()
    corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "code.jsonl", corpus)
    _query_count, query_sha = write_jsonl(SETS / "code.jsonl", queries)
    report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha}
    (SETS / "reports").mkdir(parents=True, exist_ok=True)
    (SETS / "reports" / "code.json").write_text(
        json.dumps(report, indent=1, sort_keys=True), encoding="utf-8"
    )
    print(json.dumps(report, indent=1, sort_keys=True))


if __name__ == "__main__":
    main()