Merge branch 'main' of https://github.com/rmoralespp/jsonl

Author: rmoralespp

jsonl.py

@@ -42,6 +42,48 @@
 _logger = logging.getLogger(__name__)
 _logger.addHandler(logging.NullHandler())
 
+ext_jsonl = ".jsonl"
+ext_gz = ".gz"
+ext_bz2 = ".bz2"
+ext_xz = ".xz"
+extensions = frozenset((ext_jsonl, ext_gz, ext_bz2, ext_xz))
+
+
+def _get_fileobj_extension(fileobj, /):
+    """Get the file extension based on the initial bytes of a file-like object."""
+
+    fd_position = fileobj.tell()  # Save current position
+    fileobj.seek(0)  # Go to the start of the file
+    bytes_ = fileobj.read(6)  # Read enough bytes to detect compression
+    fileobj.seek(fd_position)  # Restore the original position
+
+    if bytes_[:2] == b"\x1f\x8b":
+        # https://tools.ietf.org/html/rfc1952#page-6
+        return ext_gz
+    elif bytes_[:3] == b"\x42\x5a\x68":
+        # https://en.wikipedia.org/wiki/List_of_file_signatures
+        return ext_bz2
+    elif bytes_[:6] == b"\xfd\x37\x7a\x58\x5a\x00":
+        # https://tukaani.org/xz/xz-file-format.txt
+        return ext_xz
+    else:
+        return None
+
+
+def _get_file_extension(name, mode, /, *, fileobj=None):
+    """Get the file extension based on the filename or file-like object."""
+
+    extension = os.path.splitext(name)[1]
+    if extension in extensions:
+        return extension
+    elif mode == "rb" and fileobj:
+        return _get_fileobj_extension(fileobj)
+    elif "r" in mode:  # if not fileobj, try to open the file and detect from content
+        with open(name, "rb") as fd:
+            return _get_fileobj_extension(fd)
+    else:
+        return None
+
 
 def _looks_like_url(value, /):
     if isinstance(value, (str, urllib.request.Request)):
@@ -79,12 +121,12 @@ def _xopen(name, /, *, mode="rb", encoding=None):
 
     default = open
     openers = {
-        ".jsonl": default,
-        ".gz": gzip.open,
-        ".bz2": bz2.open,
-        ".xz": lzma.open,
+        ext_jsonl: default,
+        ext_gz: gzip.open,
+        ext_bz2: bz2.open,
+        ext_xz: lzma.open,
     }
-    extension = os.path.splitext(name)[1]
+    extension = _get_file_extension(name, mode)
     opener = openers.get(extension, default)
     return opener(name, mode=mode, encoding=encoding or _get_encoding(mode))
 
@@ -99,15 +141,15 @@ def _xfile(name, obj, /):
     :param obj: File-like an object.
     """
 
-    if name.endswith(".gz"):
+    ext = _get_file_extension(name, "rb", fileobj=obj)
+    if ext == ext_gz:
         file = gzip.GzipFile(fileobj=obj)
-    elif name.endswith(".bz2"):
+    elif ext == ext_bz2:
         file = bz2.BZ2File(obj)
-    elif name.endswith(".xz"):
+    elif ext == ext_xz:
         file = lzma.LZMAFile(obj)  # noqa: SIM115
     else:
         file = obj
-
     try:
         yield file
     finally:
@@ -344,7 +386,7 @@ def load_archive(
     :param bool broken: If true, skip broken lines (only logging a warning).
     :param Optional[Callable] json_loads: Custom function to deserialize JSON strings. By default, `json.loads` is used.
     :param Unpack[dict] json_loads_kwargs: Additional keywords to pass to `loads` of `json` provider.
-    :rtype: Generator[tuple[str, Generator[Any]]]
+    :rtype: Iterator[tuple[str, Iterator[Any]]]
     """
 
     # If a URL or Request object is provided, download the archive first.

tests/conftest.py

@@ -13,6 +13,8 @@
 import pytest
 import ujson
 
+import jsonl
+
 DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "data"))
 
 
@@ -39,6 +41,8 @@ def log_message(self, fmt, *args):  # pragma: no cover
 
     server_thread = threading.Thread(target=server.serve_forever, name="http_server")
     server_thread.start()
+    # Fix: warnings.warn(pytest.PytestUnhandledThreadExceptionWarning(msg))
+    threading.excepthook = lambda args: server.shutdown()
 
     try:
         with server.socket:
@@ -64,7 +68,7 @@ def pathlike(request):
     return request.param
 
 
-@pytest.fixture(scope="package", params=(".jsonl", ".gz", ".bz2", ".xz", ".unknown"))
+@pytest.fixture(scope="package", params=jsonl.extensions)
 def file_extension(request):
     return request.param
 

tests/test_load.py

@@ -24,8 +24,8 @@ def test_invalid_json_lines(broken):
         "[4]\n"
         "suffix\n"  # bad JSON line
     )
-    with contextlib.closing(io.StringIO(lines)) as iofile:
-        result = jsonl.load(iofile, broken=broken)
+    with contextlib.closing(io.StringIO(lines)) as fd:
+        result = jsonl.load(fd, broken=broken)
         if broken:
             assert tuple(result) == ([1, 2], [3], [4])
         else:
@@ -34,8 +34,8 @@ def test_invalid_json_lines(broken):
 
 
 def test_invalid_utf8(broken):
-    with contextlib.closing(io.BytesIO(b"\xff\xff\n[1, 2]")) as iofile:
-        result = jsonl.load(iofile, broken=broken)
+    with contextlib.closing(io.BytesIO(b"\xff\xff\n[1, 2]")) as fd:
+        result = jsonl.load(fd, broken=broken)
         if broken:
             assert tuple(result) == ([1, 2],)
         else:
@@ -78,6 +78,17 @@ def test_filepath(filepath, json_loads, pathlike):
     assert result == expected
 
 
+def test_filepath_unknown_extension_but_detected_by_signature(filepath, json_loads):
+    expected = tuple(tests.data)
+    tests.write_text(filepath, content=tests.string_data)  # Write compressed data first
+    # Rename to have an unknown extension after writing valid data
+    new_filepath = filepath + ".unknown"
+    os.rename(filepath, new_filepath)
+
+    result = tuple(jsonl.load(new_filepath, json_loads=json_loads))
+    assert result == expected
+
+
 @pytest.mark.parametrize("opener", (open, None))
 def test_filepath_using_opener(opener):
     expected = tuple(tests.data)

tests/test_load_archive.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import operator
+import os
 import shutil
 
 import pytest
@@ -14,9 +15,11 @@
     ("file1", ["file1"]),
     ("file*", ["file1", "file2"]),
 ])
+@pytest.mark.parametrize("with_unknown_ext", [True, False])
 @pytest.mark.parametrize("archive_format", ["tar", "zip"])
-def test_load_archive(pattern, match_members, tmp_dir, archive_format, file_extension):
+def test_load_archive(pattern, match_members, tmp_dir, archive_format, file_extension, with_unknown_ext):
     pattern += file_extension
+
     if archive_format == "tar":
         pattern = "./" + pattern  # tar requires a leading slash for patterns
         match_members = [f"./{member}" for member in match_members]
@@ -31,13 +34,24 @@ def test_load_archive(pattern, match_members, tmp_dir, archive_format, file_exte
             content = tests.string_data.encode(jsonl._utf_8)
             fp.write(content)
 
+    if with_unknown_ext:
+        # Rename files to have an unknown extension after writing valid data
+        for member in members:
+            new_member = member + ".unknown"
+            os.rename(root_dir / member, root_dir / new_member)
+        # Adjust pattern to match the new extension
+        pattern += ".unknown"
+
     archivepath = str(root_dir / "myarchive")
     archivepath = shutil.make_archive(archivepath, archive_format, root_dir=root_dir, base_dir=".")
 
     order_by = operator.itemgetter(0)
     expected = sorted(((name, tests.data) for name in match_members), key=order_by)
     result = jsonl.load_archive(archivepath, pattern=pattern)
     result = sorted(((name, list(data)) for name, data in result), key=order_by)
+    if with_unknown_ext:
+        # Adjust expected names to have the .unknown extension
+        expected = [(name + ".unknown", data) for name, data in expected]
     assert result == expected
 
 

tests/test_xfile.py

@@ -1,15 +1,19 @@
 # -*- coding: utf-8 -*-
+
 import bz2
 import gzip
 import io
 import lzma
 import unittest.mock
 
+import pytest
+
 import jsonl
+import tests
 
 
 def test_xfile_object(filepath):
-    obj = unittest.mock.Mock()
+    obj = unittest.mock.MagicMock()  # subscribable object
     with jsonl._xfile(filepath, obj) as result:
         if filepath.endswith(".gz"):
             assert isinstance(result, gzip.GzipFile)
@@ -36,3 +40,57 @@ def test_xfile_close(filepath):
         assert not buffer.closed
 
 
+@pytest.mark.parametrize(
+    "file_content, expected",
+    [
+        (b"\x1f\x8b\x08\x00\x00\x00", jsonl.ext_gz),  # Gzip signature
+        (b"\x42\x5a\x68\x31\x31\x39", jsonl.ext_bz2),  # Bzip2 signature
+        (b"\xfd\x37\x7a\x58\x5a\x00", jsonl.ext_xz),  # XZ signature
+        (b"\x00\x00\x00\x00\x00\x00", None),  # No matching signature
+        (b"", None),  # Empty file content
+    ],
+)
+def test_get_fileobj_extension_ok(file_content, expected):
+    fd = io.BytesIO(file_content)
+    assert jsonl._get_fileobj_extension(fd) == expected
+
+
+def test_get_fileobj_extension_restores_file_pointer():
+    file_content = b"\x1f\x8b\x08\x00\x00\x00"
+    fd = io.BytesIO(file_content)
+    initial_position = fd.tell()
+    jsonl._get_fileobj_extension(fd)
+    assert fd.tell() == initial_position
+
+
+@pytest.mark.parametrize(
+    "name, mode, fileobj, expected",
+    [
+        # Known extensions
+        ("file.jsonl", "r", None, jsonl.ext_jsonl),
+        ("file.gz", "rb", io.BytesIO(b"\x1f\x8b"), jsonl.ext_gz),
+        ("file.bz2", "rb", io.BytesIO(b"\x42\x5a\x68"), jsonl.ext_bz2),
+        ("file.xz", "rb", io.BytesIO(b"\xfd\x37\x7a\x58\x5a\x00"), jsonl.ext_xz),
+        # Unknown extensions but detected by signature
+        ("file.unknown", "rb", io.BytesIO(b"\x1f\x8b"), jsonl.ext_gz),
+        ("file.unknown", "rb", io.BytesIO(b"\x42\x5a\x68"), jsonl.ext_bz2),
+        ("file.unknown", "rb", io.BytesIO(b"\xfd\x37\x7a\x58\x5a\x00"), jsonl.ext_xz),
+        # Unknown extension but text mode
+        ("file.unknown", "r", None, None),
+        # Unknown extension but write/append binary mode
+        ("file.unknown", "wb", None, None),
+        # Unknown extensions and no fileobj or undetectable fileobj
+        ("file", "r", None, None),  # No extension
+        ("file", "rb", io.BytesIO(b""), None),  # No extension, empty fileobj
+    ],
+)
+def test_get_file_extension_ok(name, mode, fileobj, expected, tmp_dir):
+    path = tmp_dir / name
+    if not fileobj:
+        tests.write_text(path)
+    assert jsonl._get_file_extension(path, mode, fileobj=fileobj) == expected
+
+
+def test_get_file_extension_ko():
+    with pytest.raises(FileNotFoundError):
+        jsonl._get_file_extension("nonexistent.file", "r")