simtoolsz 0.2.15__tar.gz → 0.2.16__tar.gz

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/.gitignore

@@ -9,3 +9,4 @@ wheels/
 
 # venv
 .venv
+.trae/

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: simtoolsz
-Version: 0.2.15
+Version: 0.2.16
 Summary: A simple and convenient toolkit containing useful functions, classes, and methods.
 Project-URL: Homepage, https://github.com/SidneyLYZhang/simtoolsz
 Project-URL: Repository, https://github.com/SidneyLYZhang/simtoolsz.git

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "simtoolsz"
-version = "0.2.15"
+version = "0.2.16"
 description = "A simple and convenient toolkit containing useful functions, classes, and methods."
 keywords = ["tool", "collection"]
 license = { text = "MulanPSL-2.0" }

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/requirements-dev.lock

@@ -12,7 +12,7 @@
 -e file:.
 duckdb==1.4.0
     # via simtoolsz
-numpy==2.4.4
+numpy==2.5.0
     # via simtoolsz
 pendulum==3.1.0
     # via simtoolsz
@@ -20,7 +20,7 @@ polars==1.35.1
     # via simtoolsz
 polars-runtime-32==1.35.1
     # via polars
-pyarrow==23.0.1
+pyarrow==24.0.0
     # via simtoolsz
 python-dateutil==2.9.0.post0
     # via pendulum

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/requirements.lock

@@ -12,7 +12,7 @@
 -e file:.
 duckdb==1.4.0
     # via simtoolsz
-numpy==2.4.4
+numpy==2.5.0
     # via simtoolsz
 pendulum==3.1.0
     # via simtoolsz
@@ -20,7 +20,7 @@ polars==1.35.1
     # via simtoolsz
 polars-runtime-32==1.35.1
     # via polars
-pyarrow==23.0.1
+pyarrow==24.0.0
     # via simtoolsz
 python-dateutil==2.9.0.post0
     # via pendulum

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/src/simtoolsz/__init__.py

@@ -26,7 +26,7 @@ import simtoolsz.math as math
 try:
     __version__ = importlib.metadata.version("simtoolsz")
 except importlib.metadata.PackageNotFoundError:
-    __version__ = "0.2.15"
+    __version__ = "0.2.16"
 
 __author__ = "Sidney Zhang <zly@lyzhang.me>"
 

{simtoolsz-0.2.15 → simtoolsz-0.2.16}/src/simtoolsz/reader.py

@@ -108,6 +108,29 @@ def scan_tsv(filepath: Path, **kwargs) -> pl.LazyFrame:
     return read_tsv(filepath, lazy=True, **kwargs)
 
 
+_NON_ARCHIVE_EXTS = {".csv", ".tsv", ".json", ".parquet", ".ipc", ".avro", ".xlsx", ".xls", ".ods"}
+_ARCHIVE_EXTS = {".zip", ".tar", ".gz", ".bz2"}
+
+_EAGER_READERS: dict[str, Callable] = {
+    "csv": pl.read_csv,
+    "tsv": read_tsv,
+    "xls": pl.read_excel,
+    "xlsx": pl.read_excel,
+    "ods": pl.read_ods,
+    "json": pl.read_json,
+    "parquet": pl.read_parquet,
+    "ipc": pl.read_ipc,
+    "avro": pl.read_avro,
+}
+
+_LAZY_READERS: dict[str, Callable] = {
+    "csv": pl.scan_csv,
+    "parquet": pl.scan_parquet,
+    "ipc": pl.scan_ipc,
+    "tsv": scan_tsv,
+}
+
+
 def _validate_input(
     file_path: Path | str, format_type: Optional[str]
 ) -> tuple[Path, str]:
@@ -157,25 +180,7 @@ def _get_reader_mapping(lazy: bool) -> dict[str, Callable]:
     Returns:
         dict: 格式到读取函数的映射
     """
-    if lazy:
-        return {
-            "csv": pl.scan_csv,
-            "parquet": pl.scan_parquet,
-            "ipc": pl.scan_ipc,
-            "tsv": scan_tsv,
-        }
-
-    return {
-        "csv": pl.read_csv,
-        "tsv": read_tsv,
-        "xls": pl.read_excel,
-        "xlsx": pl.read_excel,
-        "ods": pl.read_ods,
-        "json": pl.read_json,
-        "parquet": pl.read_parquet,
-        "ipc": pl.read_ipc,
-        "avro": pl.read_avro,
-    }
+    return _LAZY_READERS if lazy else _EAGER_READERS
 
 
 def _handle_focus_mode(fmt: str) -> Callable:
@@ -291,6 +296,9 @@ def is_archive_file(file_path: Path) -> bool:
     Returns:
         bool: 如果是压缩文件或在压缩文件中返回True
     """
+    if file_path.suffix.lower() in _NON_ARCHIVE_EXTS:
+        if not any(p.suffix.lower() in _ARCHIVE_EXTS for p in file_path.parents):
+            return False
     return _is_archive_file(file_path) or any(
         _is_archive_file(p) for p in file_path.parents
     )
@@ -544,6 +552,8 @@ def load_data(
         >>> lazy_df = load_data("data.parquet", lazy=True)
         >>> df = load_data("data.zip/users.csv")
     """
+    file_path = Path(file_path)
+
     if is_archive_file(file_path):
         df = read_archive(file_path, format_type=format_type, **kwargs)
     else:
@@ -557,10 +567,8 @@ def load_data(
         df = reader(file_path, **kwargs)
 
     if transtype is not None:
-        if isinstance(transtype, pl.Expr):
-            df = df.with_columns(transtype)
-        if isinstance(transtype, list):
-            df = df.with_columns(*transtype)
+        exprs = [transtype] if isinstance(transtype, pl.Expr) else transtype
+        df = df.with_columns(exprs)
     return df
 
 

simtoolsz-0.2.16/tests/test_load_data_optimization.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""Tests for load_data dispatch path optimizations.
+
+Covers:
+1. is_archive_file short-circuit for known non-archive extensions
+2. str input compatibility (no AttributeError)
+3. reader mapping cached as module-level constants
+4. transtype single with_columns application (Expr / list / None)
+5. regression: archive path data.zip/users.csv still reads correctly
+"""
+
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+from unittest.mock import patch
+
+import polars as pl
+
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+from simtoolsz.reader import (
+    load_data,
+    is_archive_file,
+    _get_reader_mapping,
+    _EAGER_READERS,
+    _LAZY_READERS,
+    _NON_ARCHIVE_EXTS,
+    _ARCHIVE_EXTS,
+)
+
+
+def _make_csv(path: Path, content: str = "x,y\n1,2\n3,4\n"):
+    path.write_text(content, encoding="utf-8")
+
+
+def test_short_circuit_skips_magic_bytes_check():
+    """Known non-archive extension should not trigger is_zipfile/is_tarfile."""
+    with tempfile.TemporaryDirectory() as tmp:
+        csv_path = Path(tmp) / "data.csv"
+        _make_csv(csv_path)
+
+        with patch("simtoolsz.reader.is_zipfile") as mock_zip, \
+             patch("simtoolsz.reader.is_tarfile") as mock_tar:
+            mock_zip.return_value = False
+            mock_tar.return_value = False
+            result = is_archive_file(csv_path)
+            assert result is False
+            assert mock_zip.call_count == 0, "is_zipfile should not be called for .csv"
+            assert mock_tar.call_count == 0, "is_tarfile should not be called for .csv"
+
+
+def test_short_circuit_for_all_non_archive_exts():
+    """Every extension in _NON_ARCHIVE_EXTS should short-circuit."""
+    with tempfile.TemporaryDirectory() as tmp:
+        for ext in _NON_ARCHIVE_EXTS:
+            p = Path(tmp) / f"file{ext}"
+            _make_csv(p)
+            with patch("simtoolsz.reader.is_zipfile") as mock_zip, \
+                 patch("simtoolsz.reader.is_tarfile") as mock_tar:
+                mock_zip.return_value = False
+                mock_tar.return_value = False
+                assert is_archive_file(p) is False, f"short-circuit failed for {ext}"
+                assert mock_zip.call_count == 0, f"is_zipfile called for {ext}"
+                assert mock_tar.call_count == 0, f"is_tarfile called for {ext}"
+
+
+def test_str_input_no_attribute_error():
+    """load_data with str input should not raise AttributeError."""
+    with tempfile.TemporaryDirectory() as tmp:
+        csv_path = Path(tmp) / "data.csv"
+        _make_csv(csv_path)
+
+        df = load_data(str(csv_path))
+        assert df.shape == (2, 2)
+        assert df.columns == ["x", "y"]
+
+
+def test_reader_mapping_cached():
+    """_get_reader_mapping should return the same constant object."""
+    assert _get_reader_mapping(lazy=True) is _LAZY_READERS
+    assert _get_reader_mapping(lazy=False) is _EAGER_READERS
+    assert _get_reader_mapping(lazy=True) is _get_reader_mapping(lazy=True)
+    assert _get_reader_mapping(lazy=False) is _get_reader_mapping(lazy=False)
+
+
+def test_transtype_single_expr():
+    """transtype as single Expr applied via single with_columns."""
+    with tempfile.TemporaryDirectory() as tmp:
+        csv_path = Path(tmp) / "data.csv"
+        _make_csv(csv_path)
+
+        df = load_data(csv_path, transtype=pl.col("x").cast(pl.Int32))
+        assert df.schema["x"] == pl.Int32
+
+
+def test_transtype_list():
+    """transtype as list applied via single with_columns."""
+    with tempfile.TemporaryDirectory() as tmp:
+        csv_path = Path(tmp) / "data.csv"
+        _make_csv(csv_path)
+
+        df = load_data(csv_path, transtype=[
+            pl.col("x").cast(pl.Int32),
+            pl.col("y").cast(pl.Float64),
+        ])
+        assert df.schema["x"] == pl.Int32
+        assert df.schema["y"] == pl.Float64
+
+
+def test_transtype_none():
+    """transtype=None should not call with_columns, behavior unchanged."""
+    with tempfile.TemporaryDirectory() as tmp:
+        csv_path = Path(tmp) / "data.csv"
+        _make_csv(csv_path)
+
+        df = load_data(csv_path)
+        assert df.shape == (2, 2)
+
+
+def test_archive_path_still_reads():
+    """Regression: data.zip/users.csv should still be read correctly."""
+    with tempfile.TemporaryDirectory() as tmp:
+        zip_path = Path(tmp) / "data.zip"
+        with zipfile.ZipFile(zip_path, "w") as zf:
+            zf.writestr("users.csv", "x,y\n1,2\n3,4\n")
+
+        inner_path = zip_path / "users.csv"
+        df = load_data(inner_path)
+        assert df.shape == (2, 2)
+        assert df.columns == ["x", "y"]
+
+
+def test_unknown_extension_falls_back():
+    """Unknown extension should fall back to original magic-bytes check."""
+    with tempfile.TemporaryDirectory() as tmp:
+        no_ext_path = Path(tmp) / "data"
+        _make_csv(no_ext_path)
+
+        result = is_archive_file(no_ext_path)
+        assert result is False
+
+
+def test_constants_defined():
+    """Verify required constants exist and have expected content."""
+    assert ".csv" in _NON_ARCHIVE_EXTS
+    assert ".parquet" in _NON_ARCHIVE_EXTS
+    assert ".zip" in _ARCHIVE_EXTS
+    assert ".tar" in _ARCHIVE_EXTS
+    assert "csv" in _EAGER_READERS
+    assert "csv" in _LAZY_READERS
+
+
+def main():
+    tests = [
+        test_short_circuit_skips_magic_bytes_check,
+        test_short_circuit_for_all_non_archive_exts,
+        test_str_input_no_attribute_error,
+        test_reader_mapping_cached,
+        test_transtype_single_expr,
+        test_transtype_list,
+        test_transtype_none,
+        test_archive_path_still_reads,
+        test_unknown_extension_falls_back,
+        test_constants_defined,
+    ]
+    failed = 0
+    for t in tests:
+        try:
+            t()
+            print(f"PASS {t.__name__}")
+        except Exception as e:
+            print(f"FAIL {t.__name__}: {e}")
+            import traceback
+            traceback.print_exc()
+            failed += 1
+    if failed:
+        print(f"\n{failed}/{len(tests)} tests failed")
+        return 1
+    print(f"\n{len(tests)}/{len(tests)} tests passed")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())