library packages

.venv/lib/python3.12/site-packages/pyogrio/tests/test_util.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+from pyogrio import vsi_listtree, vsi_unlink
+from pyogrio.raw import read, write
+from pyogrio.util import vsimem_rmtree_toplevel
+
+import pytest
+
+
+def test_vsimem_rmtree_toplevel(naturalearth_lowres):
+    # Prepare test data in /vsimem/
+    meta, _, geometry, field_data = read(naturalearth_lowres)
+    meta["spatial_index"] = False
+    meta["geometry_type"] = "MultiPolygon"
+    test_dir_path = Path(f"/vsimem/test/{naturalearth_lowres.stem}.gpkg")
+    test_dir2_path = Path(f"/vsimem/test2/test2/{naturalearth_lowres.stem}.gpkg")
+
+    write(test_dir_path, geometry, field_data, **meta)
+    write(test_dir2_path, geometry, field_data, **meta)
+
+    # Check if everything was created properly with listtree
+    files = vsi_listtree("/vsimem/")
+    assert test_dir_path.as_posix() in files
+    assert test_dir2_path.as_posix() in files
+
+    # Test deleting parent dir of file in single directory
+    vsimem_rmtree_toplevel(test_dir_path)
+    files = vsi_listtree("/vsimem/")
+    assert test_dir_path.parent.as_posix() not in files
+    assert test_dir2_path.as_posix() in files
+
+    # Test deleting top-level dir of file in a subdirectory
+    vsimem_rmtree_toplevel(test_dir2_path)
+    assert test_dir2_path.as_posix() not in vsi_listtree("/vsimem/")
+
+
+def test_vsimem_rmtree_toplevel_error(naturalearth_lowres):
+    # Prepare test data in /vsimem
+    meta, _, geometry, field_data = read(naturalearth_lowres)
+    meta["spatial_index"] = False
+    meta["geometry_type"] = "MultiPolygon"
+    test_file_path = Path(f"/vsimem/pyogrio_test_{naturalearth_lowres.stem}.gpkg")
+
+    write(test_file_path, geometry, field_data, **meta)
+    assert test_file_path.as_posix() in vsi_listtree("/vsimem/")
+
+    # Deleting parent dir of non-existent file should raise an error.
+    with pytest.raises(FileNotFoundError, match="Path does not exist"):
+        vsimem_rmtree_toplevel("/vsimem/test/non-existent.gpkg")
+
+    # File should still be there
+    assert test_file_path.as_posix() in vsi_listtree("/vsimem/")
+
+    # Cleanup.
+    vsi_unlink(test_file_path)
+    assert test_file_path not in vsi_listtree("/vsimem/")