diff --git a/packages/markitdown/src/markitdown/converters/_ipynb_converter.py b/packages/markitdown/src/markitdown/converters/_ipynb_converter.py index 6bc1c1673f..62925f5ef3 100644 --- a/packages/markitdown/src/markitdown/converters/_ipynb_converter.py +++ b/packages/markitdown/src/markitdown/converters/_ipynb_converter.py @@ -24,26 +24,26 @@ def accepts( mimetype = (stream_info.mimetype or "").lower() extension = (stream_info.extension or "").lower() - if extension in ACCEPTED_FILE_EXTENSIONS: - return True - for prefix in CANDIDATE_MIME_TYPE_PREFIXES: if mimetype.startswith(prefix): - # Read further to see if it's a notebook + # Check the JSON structure, not just mentions of notebook fields. cur_pos = file_stream.tell() try: encoding = stream_info.charset or "utf-8" notebook_content = file_stream.read().decode(encoding) + notebook = json.loads(notebook_content.lstrip("\ufeff")) return ( - "nbformat" in notebook_content - and "nbformat_minor" in notebook_content + isinstance(notebook, dict) + and type(notebook.get("nbformat")) is int + and type(notebook.get("nbformat_minor")) is int + and isinstance(notebook.get("cells"), list) ) - except (ValueError, LookupError): + except (ValueError, LookupError, RecursionError): return False finally: file_stream.seek(cur_pos) - return False + return extension in ACCEPTED_FILE_EXTENSIONS def convert( self, diff --git a/packages/markitdown/tests/test_ipynb_detection.py b/packages/markitdown/tests/test_ipynb_detection.py new file mode 100644 index 0000000000..036047910e --- /dev/null +++ b/packages/markitdown/tests/test_ipynb_detection.py @@ -0,0 +1,93 @@ +import io +import json +import sys + +import pytest + +from markitdown import MarkItDown, StreamInfo +from markitdown.converters._ipynb_converter import IpynbConverter + + +@pytest.mark.parametrize( + "payload", + [ + {"description": "nbformat_minor describes a notebook version"}, + {"notebook": {"nbformat": 4, "nbformat_minor": 5, "cells": []}}, + {"nbformat": 4, "nbformat_minor": 5, "description": "Format documentation"}, + { + "nbformat": {"type": "integer"}, + "nbformat_minor": {"type": "integer"}, + "cells": [], + }, + ], + ids=["text-value", "nested-notebook", "missing-cells", "schema-fields"], +) +@pytest.mark.parametrize("entry_point", ["file", "stream"]) +def test_json_notebook_references_are_preserved(payload, entry_point, tmp_path): + """Mentioning notebook fields must not discard an ordinary JSON document.""" + text = json.dumps(payload) + converter = MarkItDown() + if entry_point == "file": + path = tmp_path / "documentation.json" + path.write_text(text, encoding="utf-8") + result = converter.convert(path) + else: + result = converter.convert_stream( + io.BytesIO(text.encode("utf-8")), + stream_info=StreamInfo(mimetype="application/json"), + ) + + assert result.markdown == text + + +@pytest.mark.parametrize( + "data, charset", + [ + (b'["nbformat", "nbformat_minor"]', "utf-8"), + (b'"nbformat_minor"', "utf-8"), + (b'{"nbformat": 4, "nbformat_minor": 5, "cells": [}', "utf-8"), + (b"[]", "unknown-charset"), + ( + b"[" * sys.getrecursionlimit() + + b'"nbformat_minor"' + + b"]" * sys.getrecursionlimit(), + "utf-8", + ), + ], + ids=["array", "string", "invalid-json", "unknown-charset", "deep-json"], +) +def test_ipynb_json_probe_rejects_non_notebooks_without_consuming_stream(data, charset): + stream = io.BytesIO(b"prefix" + data) + stream.seek(len(b"prefix")) + position = stream.tell() + + assert not IpynbConverter().accepts( + stream, StreamInfo(mimetype="application/json", charset=charset) + ) + assert stream.tell() == position + assert stream.read() == data + + +@pytest.mark.parametrize( + "encoding, charset", + [ + ("utf-8", None), + ("utf-8-sig", "utf-8"), + ("utf-8-sig", "utf-8-sig"), + ("utf-16", "utf-16"), + ], +) +def test_ipynb_json_probe_preserves_notebook_conversion(encoding, charset): + notebook = { + "nbformat": 4, + "nbformat_minor": 0, + "cells": [{"cell_type": "markdown", "source": ["# Notebook\n"]}], + } + stream = io.BytesIO(json.dumps(notebook).encode(encoding)) + stream_info = StreamInfo(mimetype="application/json", charset=charset) + + assert IpynbConverter().accepts(stream, stream_info) + assert stream.tell() == 0 + result = MarkItDown().convert_stream(stream, stream_info=stream_info) + assert result.markdown == "# Notebook\n" + assert result.title == "Notebook"