generate_report.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ast
import importlib.util
from pathlib import Path

import markdown  # type: ignore
import nbformat
from nbconvert import HTMLExporter
from weasyprint import HTML


def md_to_pdf(md_path: Path, pdf_path: Path):
    """Convert markdown file to PDF"""
    with open(md_path, encoding="utf-8") as f:
        md_text = f.read()

    html = markdown.markdown(
        md_text,
        extensions=["fenced_code", "codehilite"],
    )
    HTML(string=html, base_url=md_path.parent).write_pdf(str(pdf_path))


def ipynb_to_pdf(ipynb_path: Path, pdf_path: Path):
    """Convert Jupyter notebook to PDF via HTML"""
    with open(ipynb_path, encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=4)

    html_exporter = HTMLExporter()
    html_exporter.exclude_input_prompt = True
    html_exporter.exclude_output_prompt = True
    body, _ = html_exporter.from_notebook_node(nb)

    HTML(string=body, base_url=ipynb_path.parent).write_pdf(str(pdf_path))


def resolve_module_path(module_name):
    """Resolve module path using importlib"""
    spec = importlib.util.find_spec(module_name)
    if spec and spec.origin:
        return Path(spec.origin)
    return None


def get_matrix_calculus_imports(py_file):
    """Parse python file and extract imports from matrix_calculus"""
    with open(py_file, encoding="utf-8") as f:
        tree = ast.parse(f.read(), filename=str(py_file))

    modules = set()

    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom):
            if node.module and node.module.startswith("matrix_calculus"):
                modules.add(node.module)

        if isinstance(node, ast.Import):
            for name in node.names:
                if name.name.startswith("matrix_calculus"):
                    modules.add(name.name)

    return modules


def collect_recursive_modules(entry_file, visited=None):
    """Recursively collect matrix_calculus modules used"""
    if visited is None:
        visited = set()

    imports = get_matrix_calculus_imports(entry_file)

    files = []

    for module in imports:
        if module in visited:
            continue

        visited.add(module)
        path = resolve_module_path(module)

        if path and path.suffix == ".py":
            files.append(path)
            files.extend(collect_recursive_modules(path, visited))

    return files


def export_code(entry_file: Path, output_txt: Path):
    """Export concatenated code with matrix_calculus dependencies"""
    files = [entry_file]
    files.extend(collect_recursive_modules(entry_file))

    seen = set()
    ordered_files = []

    for f in files:
        if f not in seen:
            seen.add(f)
            ordered_files.append(f)

    with open(output_txt, "w", encoding="utf-8") as out:
        for f in ordered_files:
            out.write(f"\n========== {f.name} ==========\n")
            with open(f, encoding="utf-8") as src:
                out.write(src.read())
                out.write("\n")


def main():
    parser = argparse.ArgumentParser(description="Generate project report.")
    parser.add_argument("--description_file", "-d", required=True)
    parser.add_argument("--code", "-c", required=False, default=None)
    parser.add_argument("--output_dir", "-o", default="report")

    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    pdf_output = output_dir / "report.pdf"
    txt_output = output_dir / "code.txt"

    desc_path = Path(args.description_file)

    if desc_path.suffix == ".md":
        md_to_pdf(desc_path, pdf_output)
    elif desc_path.suffix == ".ipynb":
        ipynb_to_pdf(desc_path, pdf_output)
    else:
        raise ValueError(f"Unsupported file type: {desc_path.suffix}")

    if args.code is not None:
        code_path = Path(args.code)

        export_code(code_path, txt_output)


if __name__ == "__main__":
    main()