fix file read error

app.py

@@ -8,170 +8,200 @@ from typing import Any, List, Tuple, Dict, Union
 import gradio as gr
 
 
-# -----------------------------
-# Safe unpickling (no imports)
-# -----------------------------
+# =========================
+# Safe unpickling machinery
+# =========================
 class _Record(dict):
-    """Placeholder for any external class; __setstate__ writes into this dict."""
+    """
+    Placeholder for any external class instance encountered during unpickling.
+    We implement __setstate__ so state can be applied into this dict safely.
+    """
     def __setstate__(self, state):
         if isinstance(state, dict):
             self.update(state)
         else:
-            # keep non-dict state somewhere visible
+            # Keep non-dict states visible rather than failing
             self["_state"] = state
 
+
 class _PlaceholderType(type):
+    """
+    A metaclass that returns _Record() when the placeholder 'type' is called.
+    This handles REDUCE/callable paths in pickle protocols.
+    """
     def __call__(cls, *args, **kwargs):
-        # Called for REDUCE/callable cases
         return _Record()
 
+
 class _Placeholder(metaclass=_PlaceholderType):
+    """
+    A 'type' whose __new__ also returns _Record(), satisfying NEWOBJ/NEWOBJ_EX.
+    """
     def __new__(cls, *args, **kwargs):
-        # Used by NEWOBJ/NEWOBJ_EX
         return _Record()
 
+
 class _SafeUnpickler(pickle.Unpickler):
-    """Never import external modules; map everything to a benign placeholder type."""
+    """
+    Unpickler that never imports external modules/classes.
+    Every GLOBAL reference becomes _Placeholder, which yields _Record objects.
+    """
     def find_class(self, module, name):
-        return _Placeholder  # returning a 'type' satisfies NEWOBJ paths
+        # Never import arbitrary classes; map to our safe placeholder 'type'
+        return _Placeholder
 
 
 def _to_builtin(obj: Any) -> Any:
-    """Recursively convert placeholders/bytes/sets into JSON-friendly Python builtins."""
+    """
+    Recursively convert placeholders and non-JSON-friendly types into
+    plain Python builtins (dict/list/str/ints), suitable for JSON dumping.
+    """
     if isinstance(obj, _Record) or isinstance(obj, dict):
         return {_to_builtin(k): _to_builtin(v) for k, v in obj.items()}
     if isinstance(obj, (list, tuple, set)):
         return [_to_builtin(x) for x in obj]
     if isinstance(obj, bytes):
+        # Try decoding as UTF-8; otherwise keep a small descriptor
         try:
             return obj.decode("utf-8")
         except Exception:
-            # keep as repr if not decodable
             return {"__bytes__": True, "len": len(obj)}
     return obj
 
 
 def safe_unpickle(data: bytes) -> Any:
-    """Safe, single-object unpickle to builtins."""
+    """
+    Safely unpickle a single pickle object from bytes without importing modules.
+    """
     obj = _SafeUnpickler(io.BytesIO(data)).load()
     return _to_builtin(obj)
 
 
-# ---------------------------------------
-# Support multi-object pickle concatenation
-# ---------------------------------------
-PROTO_MARKER = re.compile(br"\x80[\x03-\x05]")  # common proto 3/4/5
+# =========================
+# Multi-pickle file support
+# =========================
+# Common pickle protocol markers: 0x80 0x03/0x04/0x05
+PROTO_MARKER = re.compile(br"\x80[\x03-\x05]")
+
 
 def split_pickle_stream(data: bytes) -> List[Tuple[int, int, bytes]]:
     """
-    Some tools append multiple pickles into one file. Split by protocol markers.
-    If none found, return the whole file as a single segment.
+    Some tools concatenate multiple pickles in one file.
+    Split the byte stream by protocol markers.
+    Returns a list of (start, end, segment_bytes).
     """
     idxs = [m.start() for m in PROTO_MARKER.finditer(data)]
     if not idxs:
+        # Not a concatenated file (or older protocol); treat whole file as one segment
         return [(0, len(data), data)]
     idxs.append(len(data))
-    segs = []
+    segments = []
     for i in range(len(idxs) - 1):
         s, e = idxs[i], idxs[i + 1]
-        segs.append((s, e, data[s:e]))
-    return segs
+        segments.append((s, e, data[s:e]))
+    return segments
 
 
-# -----------------------------
-# Gradio logic
-# -----------------------------
+# =========================
+# Gradio handlers / UI
+# =========================
 def handle_file(
-    file: gr.File | None,
-    pretty: bool,
-    limit_preview: int,
-) -> Tuple[str, Dict[str, Any] | List[Any] | None, str | None]:
+    file: Union[bytes, io.BufferedReader, None],
+    pretty_json: bool,
+    max_error_lines: int,
+) -> Tuple[str, Union[Dict[str, Any], List[Any], None], Union[str, None]]:
     """
-    Returns: (summary_text, json_preview, download_path)
+    Gradio callback:
+    - 'file' can be bytes (when File(type="binary")) or a file-like object.
+    - Returns (summary_text, json_preview_object, download_json_path).
     """
     if file is None:
         return "Please upload a .pkl (pickle) file.", None, None
 
-    raw = file.read()
+    # Accept both raw bytes and file-like objects
+    if isinstance(file, (bytes, bytearray)):
+        raw = bytes(file)
+    else:
+        raw = file.read()
+
     size_kb = len(raw) / 1024.0
 
+    # Split and safely unpickle each segment
     segments = split_pickle_stream(raw)
     results: List[Any] = []
     errors: List[str] = []
-
     for i, (_, _, seg) in enumerate(segments, 1):
         try:
             results.append(safe_unpickle(seg))
         except Exception as e:
             errors.append(f"Segment {i}: {type(e).__name__}: {e}")
 
-    # If it's a single segment, show the object directly; otherwise show a list
+    # Single segment -> show the object; multiple -> show a list
     output_obj: Union[Dict[str, Any], List[Any]]
     if len(results) == 1:
         output_obj = results[0]
     else:
         output_obj = results
 
-    # Build summary text
+    # Build summary
     summary_lines = [
         f"File size: {size_kb:.1f} KB",
         f"Detected pickle segments: {len(segments)}",
     ]
     if errors:
         summary_lines.append(f"Errors: {len(errors)}")
-        summary_lines.extend(errors[:limit_preview])
-
+        summary_lines.extend(errors[:max_error_lines])
     summary = "\n".join(summary_lines)
 
-    # Create a temporary JSON file for download
-    tmp = tempfile.NamedTemporaryFile(prefix="pickle_view_", suffix=".json", delete=False)
+    # Create a temporary JSON for download
+    tmp = tempfile.NamedTemporaryFile(prefix="depickle_", suffix=".json", delete=False)
     with open(tmp.name, "w", encoding="utf-8") as f:
-        if pretty:
+        if pretty_json:
             json.dump(output_obj, f, ensure_ascii=False, indent=2)
         else:
             json.dump(output_obj, f, ensure_ascii=False, separators=(",", ":"))
     download_path = tmp.name
 
-    # For the on-page JSON preview, Gradio's JSON component handles dict/list nicely
-    # (It doesn't use the 'pretty' toggle, which only affects the downloadable file.)
-    json_preview = output_obj
+    # For the on-page preview, Gradio's JSON component accepts dict/list directly
+    return summary, output_obj, download_path
 
-    return summary, json_preview, download_path
 
-
-with gr.Blocks(title="Pickle Viewer (Safe)") as demo:
+with gr.Blocks(title="DePickle — Safe Pickle Viewer") as demo:
     gr.Markdown(
         """
-# 🥒 Pickle Viewer (Safe)
-Safely deserialize Python pickle files **without executing arbitrary code**, then preview and download as JSON.
+# 🥒 DePickle — Safe Pickle Viewer
+
+Safely deserialize Python pickle files **without executing arbitrary code** and preview as JSON.
 
 - **Safe**: Unknown classes are mapped to harmless placeholders (no imports run).
-- **Multi-pickle**: Supports files that contain multiple concatenated pickle objects.
-- **Download**: Get the parsed JSON as a file.
+- **Multi-pickle**: Supports files containing multiple concatenated pickle objects.
+- **Download**: Export the parsed content as a JSON file.
 
-> Tip: This tool is for **inspection** only. If you need to *execute* real classes, do it locally and only on trusted data.
+> Note: This tool is for **inspection** only. If you need to execute real classes, do it locally and only on trusted data.
 """
     )
 
     with gr.Row():
-        file_in = gr.File(label="Upload .pkl file", file_types=[".pkl", ".pickle"], type="binary")
+        file_in = gr.File(
+            label="Upload .pkl / .pickle",
+            file_types=[".pkl", ".pickle"],
+            type="binary",  # returns bytes; our handler supports bytes/file
+        )
         with gr.Column():
             pretty = gr.Checkbox(value=True, label="Pretty JSON (for download)")
-            limit_preview = gr.Slider(
-                1, 50, value=10, step=1,
-                label="Max error lines to show in summary"
-            )
+            max_lines = gr.Slider(1, 50, value=10, step=1, label="Max error lines in summary")
 
     run_btn = gr.Button("Deserialize Safely")
-    summary_out = gr.Textbox(label="Summary", lines=6)
+    summary_out = gr.Textbox(label="Summary", lines=8)
     json_out = gr.JSON(label="Parsed JSON Preview")
     download_out = gr.File(label="Download JSON")
 
     run_btn.click(
         fn=handle_file,
-        inputs=[file_in, pretty, limit_preview],
+        inputs=[file_in, pretty, max_lines],
         outputs=[summary_out, json_out, download_out],
     )
 
 if __name__ == "__main__":
+    # In Spaces, Gradio will be launched by the runtime; locally you can run it yourself.
     demo.launch()