extract code blocks as files from GLM5.2 chat -> into zip archive

mdextract.py

#!/usr/bin/env python3
"""
mdextract — Extract fenced code blocks from Markdown into a ZIP archive or directory.

It looks for headers of the form:

    ### `relative/path.ext`

followed by a fenced block:

    ```lang
    code here
    ```

and writes the discovered code to either a .zip archive or a directory tree.
"""

from __future__ import annotations

import argparse
import logging
import re
import sys
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Iterator

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# Matches: ### `path/to/file.ext`  \n  ```lang \n <code> \n ```
# - filename may contain word chars, dots, slashes, hyphens
# - code body is matched non-greedily so backticks inside the code are tolerated
BLOCK_RE = re.compile(
    r"### `([\w./\-]+)`\n+```(\w*)\n(.*?)\n```",
    re.DOTALL,
)

log = logging.getLogger("mdextract")


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass
class CodeBlock:
    filename: str
    language: str
    code: str

    def safe_path(self) -> Path:
        """Return a Path with any traversal / absolute segments removed."""
        parts = [p for p in self.filename.replace("\\", "/").split("/")
                 if p not in ("", ".", "..")]
        return Path(*parts) if parts else Path(f"untitled.{self.language or 'txt'}")


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------

def parse_markdown(text: str) -> Iterator[CodeBlock]:
    """Yield CodeBlock objects from a Markdown string."""
    for i, (filename, lang, code) in enumerate(BLOCK_RE.findall(text)):
        if not filename:
            filename = f"extracted_code_{i}.{lang or 'txt'}"
        yield CodeBlock(filename=filename, language=lang, code=code.strip())


def read_input(paths: list[str]) -> str:
    """Read text from one or more files (use '-' for stdin)."""
    chunks: list[str] = []
    for p in paths:
        if p == "-":
            chunks.append(sys.stdin.read())
            continue
        path = Path(p)
        if not path.is_file():
            raise FileNotFoundError(f"Input file not found: {path}")
        chunks.append(path.read_text(encoding="utf-8"))
    return "\n".join(chunks)


def write_to_zip(blocks: Iterable[CodeBlock], archive: Path, force: bool) -> int:
    if archive.exists() and not force:
        raise FileExistsError(
            f"Output archive already exists: {archive} (use --force to overwrite)"
        )
    count = 0
    with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
        for block in blocks:
            zf.writestr(str(block.safe_path()), block.code)
            log.info("added %s", block.filename)
            count += 1
    return count


def write_to_dir(blocks: Iterable[CodeBlock], out_dir: Path, force: bool) -> int:
    count = 0
    for block in blocks:
        target = out_dir / block.safe_path()
        if target.exists() and not force:
            raise FileExistsError(
                f"File exists (use --force to overwrite): {target}"
            )
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(block.code, encoding="utf-8")
        log.info("wrote %s", target)
        count += 1
    return count


def list_blocks(blocks: Iterable[CodeBlock]) -> int:
    count = 0
    for block in blocks:
        print(f"{block.filename}\t[{block.language}]")
        count += 1
    return count


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="mdextract",
        description="Extract fenced code blocks (with ### `filename` headers) "
                    "from Markdown files into a ZIP archive or directory.",
        epilog=(
            "Examples:\n"
            "  mdextract notes.md -o code.zip\n"
            "  mdextract notes.md --to-dir ./src --force\n"
            "  mdextract a.md b.md c.md -o combined.zip\n"
            "  cat notes.md | mdextract - -o code.zip\n"
            "  mdextract notes.md --list\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("inputs", nargs="+",
                   help="One or more Markdown files (use '-' to read from stdin).")
    p.add_argument("-o", "--output", default="",
                   help="Output ZIP archive name (default: %(default)s). "
                        "Ignored when --to-dir is given.")
    p.add_argument("--to-dir", metavar="DIR",
                   help="Write extracted files to a directory tree instead of a ZIP archive.")
    p.add_argument("-l", "--list", action="store_true",
                   help="List discovered code blocks and exit without writing anything.")
    p.add_argument("-n", "--dry-run", action="store_true",
                   help="Show what would be written, but do not write any files.")
    p.add_argument("-f", "--force", action="store_true",
                   help="Overwrite existing output files / archive.")
    p.add_argument("-v", "--verbose", action="count", default=0,
                   help="Increase log verbosity (-v, -vv).")
    p.add_argument("-q", "--quiet", action="count", default=0,
                   help="Decrease log verbosity (-q, -qq).")
    p.add_argument("--version", action="version", version="mdextract 1.0.0")
    return p


def configure_logging(verbose: int, quiet: int) -> None:
    level = logging.WARNING + (quiet * 10) - (verbose * 10)
    level = max(level, logging.DEBUG)
    logging.basicConfig(
        level=level,
        format="%(levelname)s: %(message)s",
        stream=sys.stderr,
    )


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    configure_logging(args.verbose, args.quiet)

    # ---- read input -------------------------------------------------------
    try:
        text = read_input(args.inputs)
    except FileNotFoundError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    except OSError as exc:
        print(f"error: could not read input: {exc}", file=sys.stderr)
        return 2

    # ---- output filename default ------------------------------------------
    if not args.output:
        input_filename = args.inputs[0]
        out_filename = f"extracted_code_for_{input_filename}.zip"
    else:
        out_filename = arg.output 


    # ---- parse ------------------------------------------------------------
    blocks = list(parse_markdown(text))
    if not blocks:
        print("warning: no code blocks found in input.", file=sys.stderr)
        return 1
    log.debug("discovered %d code block(s)", len(blocks))

    # ---- list mode --------------------------------------------------------
    if args.list:
        return 0 if list_blocks(blocks) else 1

    # ---- dry-run mode -----------------------------------------------------
    if args.dry_run:
        for b in blocks:
            print(f"would write: {b.safe_path()} ({len(b.code)} bytes)")
        print(f"dry run: {len(blocks)} file(s) would be written.", file=sys.stderr)
        return 0

    # ---- write ------------------------------------------------------------
    try:
        if args.to_dir:
            out_dir = Path(args.to_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
            n = write_to_dir(blocks, out_dir, args.force)
            print(f"Wrote {n} file(s) to {out_dir}/", file=sys.stderr)
        else:
            archive = Path(out_filename)
            n = write_to_zip(blocks, archive, args.force)
            print(f"Wrote {n} file(s) to {archive}", file=sys.stderr)
    except FileExistsError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 3
    except OSError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 4

    return 0


if __name__ == "__main__":
    sys.exit(main())