thepowerfuldeez/the-stack-v2-train-smol-ids-updated

from collections.abc import Iterable
from dataclasses import dataclass
from html import unescape
from pathlib import Path

import regex as re

ALLOWED_LANGUAGES = {
    "Ant Build System",
    "AsciiDoc",
    "C",
    "C#",
    "C++",
    "CMake",
    "Dockerfile",
    "Go",
    "Go Module",
    "Gradle",
    "Groovy",
    "HTML",
    "INI",
    "Java",
    "Java Properties",
    "JavaScript",
    "JSON",
    "JSON with Comments",
    "Kotlin",
    "Lua",
    "M4Sugar",
    "Makefile",
    "Markdown",
    "Maven POM",
    "PHP",
    "Python",
    "R",
    "RDoc",
    "reStructuredText",
    "RMarkdown",
    "Ruby",
    "Rust",
    "Shell",
    "SQL",
    "Swift",
    "Text",
    "TOML",
    "TypeScript",
    "YAML",
}

AUTOGEN_PHRASES = (
    "auto-generated",
    "autogenerated",
    "automatically generated",
    "generated automatically",
    "this file is generated",
)

RE_BASE64 = re.compile(r"[a-zA-Z0-9+/\n=]{64,}")
RE_HEXSEQ = re.compile(r"(?:\b(?:0x|\\x)?[0-9a-fA-F]{2}(?:,|\b\s*)){8,}")
RE_UNICODE = re.compile(r"(?:\\u[0-9a-fA-F]{4}){8,}")


@dataclass(slots=True)
class FileStats:
    path: str
    size: int
    nlines: int
    avg_line_len: float
    max_line_len: int
    alpha_ratio: float
    is_binary: bool
    language: str | None


EXT_TO_LANGUAGE = {
    # minimal mapping; can be extended, or replaced by enry if available
    ".c": "C",
    ".h": "C",
    ".cc": "C++",
    ".cpp": "C++",
    ".hpp": "C++",
    ".cs": "C#",
    ".cmake": "CMake",
    "CMakeLists.txt": "CMake",
    "Dockerfile": "Dockerfile",
    ".go": "Go",
    ".mod": "Go Module",
    ".gradle": "Gradle",
    ".groovy": "Groovy",
    ".html": "HTML",
    ".ini": "INI",
    ".java": "Java",
    ".properties": "Java Properties",
    ".js": "JavaScript",
    ".mjs": "JavaScript",
    ".ts": "TypeScript",
    ".json": "JSON",
    ".jsonc": "JSON with Comments",
    ".kt": "Kotlin",
    ".lua": "Lua",
    ".m4": "M4Sugar",
    "Makefile": "Makefile",
    ".md": "Markdown",
    "pom.xml": "Maven POM",
    ".php": "PHP",
    ".py": "Python",
    ".r": "R",
    ".rdoc": "RDoc",
    ".rst": "reStructuredText",
    ".rmd": "RMarkdown",
    ".rb": "Ruby",
    ".rs": "Rust",
    ".sh": "Shell",
    ".sql": "SQL",
    ".swift": "Swift",
    ".txt": "Text",
    ".toml": "TOML",
    ".yaml": "YAML",
    ".yml": "YAML",
}


def detect_language(path: str) -> str | None:
    name = Path(path).name
    if name in EXT_TO_LANGUAGE:
        return EXT_TO_LANGUAGE[name]
    ext = Path(path).suffix
    return EXT_TO_LANGUAGE.get(ext)


def is_generated_content(text: str, language: str | None) -> bool:
    # Primary: phrase scan in first 5 lines
    head = "\n".join(text.splitlines()[:5]).lower()
    if any(p in head for p in AUTOGEN_PHRASES):
        return True
    # Optional: try enry if available
    try:
        import enry  # type: ignore

        return bool(enry.is_generated(text))
    except Exception:
        return False


def visible_text_from_html(text: str) -> str:
    # Very light-weight visibility heuristic: strip tags, scripts/styles
    text = re.sub(r"<script[\s\S]*?</script>", " ", text)
    text = re.sub(r"<style[\s\S]*?</style>", " ", text)
    text = re.sub(r"<!--.*?-->", " ", text)
    text = re.sub(r"<[^>]+>", " ", text)
    text = unescape(text)
    return re.sub(r"\s+", " ", text).strip()


def compute_stats(text: str, path: str) -> FileStats:
    lines = text.splitlines() or [""]
    nlines = len(lines)
    lengths = [len(line) for line in lines]
    avg_len = sum(lengths) / max(1, nlines)
    max_len = max(lengths) if lengths else 0
    alpha = sum(c.isalpha() for c in text)
    alpha_ratio = alpha / max(1, len(text))
    # naive binary detection
    is_binary = "\x00" in text or (sum(1 for c in text if ord(c) < 9) > 0)
    language = detect_language(path)
    return FileStats(
        path=path,
        size=len(text.encode("utf-8", "ignore")),
        nlines=nlines,
        avg_line_len=avg_len,
        max_line_len=max_len,
        alpha_ratio=alpha_ratio,
        is_binary=is_binary,
        language=language,
    )


def encoded_data_fraction(text: str) -> tuple[int, int]:
    matched = 0
    longest = 0
    for rx in (RE_BASE64, RE_HEXSEQ, RE_UNICODE):
        for m in rx.finditer(text):
            seg_len = len(m.group(0))
            matched += seg_len
            longest = max(longest, seg_len)
    return matched, longest


def _binary_or_size_check(stats: FileStats) -> bool:
    if stats.is_binary or stats.size == 0 or stats.size > 10 * 1024 * 1024:
        return False
    return True


def _length_checks(stats: FileStats) -> bool:
    long_langs = {"HTML", "JSON", "Markdown", "Roff", "Roff Manpage", "SMT", "TeX", "Text", "XML"}
    if stats.language in long_langs:
        return stats.max_line_len <= 100_000
    if stats.nlines > 100_000:
        return False
    if stats.avg_line_len > 100:
        return False
    if stats.max_line_len > 1_000:
        return False
    return True


def _autogen_check(text: str, language: str | None) -> bool:
    return not is_generated_content(text, language)


def _alpha_check(stats: FileStats) -> bool:
    if stats.language in {"Motorola 68K Assembly", "WebAssembly"}:
        return True
    return stats.alpha_ratio >= 0.25


def _encoded_check(text: str) -> bool:
    matched, longest = encoded_data_fraction(text)
    if longest > 1024:
        return False
    if matched / max(1, len(text)) > 0.5:
        return False
    return True


def basic_filters(text: str, stats: FileStats) -> bool:
    return (
        _binary_or_size_check(stats)
        and _length_checks(stats)
        and _autogen_check(text, stats.language)
        and _alpha_check(stats)
        and _encoded_check(text)
    )


def language_specific_filters(text: str, stats: FileStats) -> bool:
    lang = stats.language
    if not lang:
        return False
    # Language allowlist
    if lang not in ALLOWED_LANGUAGES:
        return False

    if lang in {"Text", "JSON", "YAML", "Web Ontology Language", "Graphviz (DOT)"}:
        if stats.nlines > 512:
            return False

    if lang == "HTML":
        visible = visible_text_from_html(text)
        if len(visible) < 100:
            return False
        if len(visible) / max(1, len(text)) < 0.2:
            return False

    if lang == "Text":
        name = Path(stats.path).name.lower()
        base = Path(stats.path).stem.lower()
        if ("requirement" not in name) and (
            base not in {"readme", "notes", "todo", "description", "cmakelists"}
        ):
            return False

    return True


def should_keep_file(path: str, text: str, *, languages: Iterable[str] | None = None) -> bool:
    stats = compute_stats(text, path)
    if languages is not None:
        if stats.language not in set(languages):
            return False
    return basic_filters(text, stats) and language_specific_filters(text, stats)