pdf reader

pdf_index/app.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+from pathlib import Path
+
+import fitz  # PyMuPDF
+import streamlit as st
+from streamlit_pdf_viewer import pdf_viewer
+
+
+@dataclass
+class TocItem:
+    level: int
+    title: str
+    page_from: int  # 1-based
+    page_to: Optional[int]  # 1-based inclusive; None means until end
+
+
+def read_toc(doc: fitz.Document) -> List[Tuple[int, str, int]]:
+    # Returns list of (level, title, page) where page is 1-based per PyMuPDF
+    toc: List[Tuple[int, str, int]] = []
+    try:
+        get_toc: Any = getattr(doc, "get_toc", None)
+        if callable(get_toc):
+            toc = get_toc(simple=True)  # type: ignore[no-any-return]
+    except Exception:
+        toc = []
+    return [(lvl, title, max(1, pg)) for (lvl, title, pg) in toc]
+
+
+def normalize_ranges(toc: List[Tuple[int, str, int]], page_count: int) -> List[TocItem]:
+    if not toc:
+        return []
+    items: List[TocItem] = []
+    for i, (lvl, title, page) in enumerate(toc):
+        start = min(max(1, page), page_count)
+        if i + 1 < len(toc):
+            next_page = toc[i + 1][2]
+            end = max(1, min(page_count, next_page - 1))
+            if end < start:
+                end = start
+        else:
+            end = page_count
+        items.append(TocItem(level=lvl, title=title, page_from=start, page_to=end))
+    return items
+
+
+def _hash_doc(doc:fitz.Document):
+    "This is a fake hash, ENSURE GLOBAL DOCUMENT SAME"
+    return "12"
+
+@st.cache_resource(hash_funcs={fitz.Document:_hash_doc})
+def slice_pdf_pages(src_doc: fitz.Document, page_from: int, page_to: int) -> bytes:
+    # Create a new PDF with selected 1-based inclusive page range
+    new_pdf = fitz.open()
+    try:
+        start_i = max(1, page_from) - 1
+        end_i = max(start_i, page_to - 1)
+        for p in range(start_i, min(end_i, src_doc.page_count - 1) + 1):
+            new_pdf.insert_pdf(src_doc, from_page=p, to_page=p)
+        out = new_pdf.tobytes()
+        return out
+    finally:
+        new_pdf.close()
+        src_doc.close()
+
+
+def format_label(item: TocItem) -> str:
+    return f"{item.title} ({item.page_from:03d} - {item.page_to:03d})"
+
+@st.cache_resource
+def read_pdf():
+    pdf_path = Path("doc.pdf")
+    if not pdf_path.exists():
+        st.error("找不到doc.pdf")
+        st.stop()
+    
+    doc = fitz.open(pdf_path, filetype="pdf")
+    def _close_doc():
+        "Never close doc due to cache in global streamlit app."
+        pass
+    doc.close = _close_doc
+    page_count = doc.page_count
+
+    # 读取目录
+    raw_toc = read_toc(doc)
+    items = normalize_ranges(raw_toc, page_count)
+        
+    return doc,page_count,items
+
+
+def main():
+    st.set_page_config(page_title="PDF 目录查看器", layout="wide")
+    st.title("PDF 目录查看器")
+
+    doc,page_count,items = read_pdf()
+
+    st.subheader("目录")
+    labels = ["请选择"] + [format_label(it) for it in items]
+    selection = st.selectbox("选择章节", labels, index=0)
+    if selection == "请选择":
+        st.stop()
+    idx = labels.index(selection)
+    chosen = items[idx]
+    selected_range = (chosen.page_from, chosen.page_to or page_count)
+
+    rng_from, rng_to = selected_range
+    
+    st.subheader("预览")
+    try:
+        sliced_bytes = slice_pdf_pages(doc, rng_from, rng_to)
+        st.download_button("下载",sliced_bytes,file_name=f"{chosen.title}.pdf")
+        pdf_viewer(sliced_bytes,rendering=st.session_state.get("render","unwrap"))
+        # st.pdf(io.BytesIO(sliced_bytes), height=height, key="pdf_preview")
+    except Exception as e:
+        st.error(f"渲染失败：{e}")
+        
+    st.selectbox("使用其他渲染方式",["unwrap","legacy_embed","legacy_iframe"],key="render")
+
+
+if __name__ == "__main__":
+    main()