pyogrio 0.9.0__cp312-cp312-manylinux_2_28_aarch64.whl → 0.10.0__cp312-cp312-manylinux_2_28_aarch64.whl

pyogrio/tests/test_core.py

@@ -1,29 +1,35 @@
+from pathlib import Path
+
 import numpy as np
-from numpy import array_equal, allclose
-import pytest
+from numpy import allclose, array_equal
 
 from pyogrio import (
-    __gdal_version__,
     __gdal_geos_version__,
+    __gdal_version__,
+    detect_write_driver,
+    get_gdal_config_option,
+    get_gdal_data_path,
     list_drivers,
     list_layers,
     read_bounds,
     read_info,
     set_gdal_config_options,
-    get_gdal_config_option,
-    get_gdal_data_path,
+    vsi_listtree,
+    vsi_rmtree,
+    vsi_unlink,
 )
-from pyogrio.core import detect_write_driver
 from pyogrio._compat import GDAL_GE_38
-from pyogrio.errors import DataSourceError, DataLayerError
-from pyogrio.tests.conftest import HAS_SHAPELY, prepare_testfile
-
 from pyogrio._env import GDALEnv
+from pyogrio.errors import DataLayerError, DataSourceError
+from pyogrio.raw import read, write
+from pyogrio.tests.conftest import START_FID, prepare_testfile, requires_shapely
+
+import pytest
 
 with GDALEnv():
     # NOTE: this must be AFTER above imports, which init the GDAL and PROJ data
     # search paths
-    from pyogrio._ogr import ogr_driver_supports_write, has_gdal_data, has_proj_data
+    from pyogrio._ogr import has_gdal_data, has_proj_data, ogr_driver_supports_write
 
 
 try:
@@ -151,7 +157,16 @@ def test_list_drivers():
     assert len(drivers) == len(expected)
 
 
-def test_list_layers(naturalearth_lowres, naturalearth_lowres_vsi, test_fgdb_vsi):
+def test_list_layers(
+    naturalearth_lowres,
+    naturalearth_lowres_vsi,
+    naturalearth_lowres_vsimem,
+    line_zm_file,
+    curve_file,
+    curve_polygon_file,
+    multisurface_file,
+    no_geometry_file,
+):
     assert array_equal(
         list_layers(naturalearth_lowres), [["naturalearth_lowres", "Polygon"]]
     )
@@ -160,21 +175,27 @@ def test_list_layers(naturalearth_lowres, naturalearth_lowres_vsi, test_fgdb_vsi
         list_layers(naturalearth_lowres_vsi[1]), [["naturalearth_lowres", "Polygon"]]
     )
 
+    assert array_equal(
+        list_layers(naturalearth_lowres_vsimem),
+        [["naturalearth_lowres", "MultiPolygon"]],
+    )
+
     # Measured 3D is downgraded to plain 3D during read
     # Make sure this warning is raised
     with pytest.warns(
         UserWarning, match=r"Measured \(M\) geometry types are not supported"
     ):
-        fgdb_layers = list_layers(test_fgdb_vsi)
-        # GDAL >= 3.4.0 includes 'another_relationship' layer
-        assert len(fgdb_layers) >= 7
+        assert array_equal(list_layers(line_zm_file), [["line_zm", "LineString Z"]])
 
-        # Make sure that nonspatial layer has None for geometry
-        assert array_equal(fgdb_layers[0], ["basetable_2", None])
+    # Curve / surface types are downgraded to plain types
+    assert array_equal(list_layers(curve_file), [["curve", "LineString"]])
+    assert array_equal(list_layers(curve_polygon_file), [["curvepolygon", "Polygon"]])
+    assert array_equal(
+        list_layers(multisurface_file), [["multisurface", "MultiPolygon"]]
+    )
 
-        # Confirm that measured 3D is downgraded to plain 3D during read
-        assert array_equal(fgdb_layers[3], ["test_lines", "MultiLineString Z"])
-        assert array_equal(fgdb_layers[6], ["test_areas", "MultiPolygon Z"])
+    # Make sure that nonspatial layer has None for geometry
+    assert array_equal(list_layers(no_geometry_file), [["no_geometry", None]])
 
 
 def test_list_layers_bytes(geojson_bytes):
@@ -184,6 +205,13 @@ def test_list_layers_bytes(geojson_bytes):
     assert layers[0, 0] == "test"
 
 
+def test_list_layers_nonseekable_bytes(nonseekable_bytes):
+    layers = list_layers(nonseekable_bytes)
+
+    assert layers.shape == (1, 2)
+    assert layers[0, 1] == "Point"
+
+
 def test_list_layers_filelike(geojson_filelike):
     layers = list_layers(geojson_filelike)
 
@@ -191,22 +219,18 @@ def test_list_layers_filelike(geojson_filelike):
     assert layers[0, 0] == "test"
 
 
-def test_read_bounds(naturalearth_lowres):
-    fids, bounds = read_bounds(naturalearth_lowres)
-    assert fids.shape == (177,)
-    assert bounds.shape == (4, 177)
-
-    assert fids[0] == 0
-    # Fiji; wraps antimeridian
-    assert allclose(bounds[:, 0], [-180.0, -18.28799, 180.0, -16.02088])
-
+@pytest.mark.parametrize(
+    "testfile",
+    ["naturalearth_lowres", "naturalearth_lowres_vsimem", "naturalearth_lowres_vsi"],
+)
+def test_read_bounds(testfile, request):
+    path = request.getfixturevalue(testfile)
+    path = path if not isinstance(path, tuple) else path[1]
 
-def test_read_bounds_vsi(naturalearth_lowres_vsi):
-    fids, bounds = read_bounds(naturalearth_lowres_vsi[1])
+    fids, bounds = read_bounds(path)
     assert fids.shape == (177,)
     assert bounds.shape == (4, 177)
-
-    assert fids[0] == 0
+    assert fids[0] == START_FID[Path(path).suffix]
     # Fiji; wraps antimeridian
     assert allclose(bounds[:, 0], [-180.0, -18.28799, 180.0, -16.02088])
 
@@ -218,6 +242,13 @@ def test_read_bounds_bytes(geojson_bytes):
     assert allclose(bounds[:, 0], [-180.0, -18.28799, 180.0, -16.02088])
 
 
+def test_read_bounds_nonseekable_bytes(nonseekable_bytes):
+    fids, bounds = read_bounds(nonseekable_bytes)
+    assert fids.shape == (1,)
+    assert bounds.shape == (4, 1)
+    assert allclose(bounds[:, 0], [1, 1, 1, 1])
+
+
 def test_read_bounds_filelike(geojson_filelike):
     fids, bounds = read_bounds(geojson_filelike)
     assert fids.shape == (3,)
@@ -285,12 +316,9 @@ def test_read_bounds_bbox(naturalearth_lowres_all_ext):
     fids, bounds = read_bounds(naturalearth_lowres_all_ext, bbox=(-85, 8, -80, 10))
 
     assert fids.shape == (2,)
-    if naturalearth_lowres_all_ext.suffix == ".gpkg":
-        # fid in gpkg is 1-based
-        assert array_equal(fids, [34, 35])  # PAN, CRI
-    else:
-        # fid in other formats is 0-based
-        assert array_equal(fids, [33, 34])  # PAN, CRI
+    fids_expected = np.array([33, 34])  # PAN, CRI
+    fids_expected += START_FID[naturalearth_lowres_all_ext.suffix]
+    assert array_equal(fids, fids_expected)
 
     assert bounds.shape == (4, 2)
     assert allclose(
@@ -302,9 +330,7 @@ def test_read_bounds_bbox(naturalearth_lowres_all_ext):
     )
 
 
-@pytest.mark.skipif(
-    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
-)
+@requires_shapely
 @pytest.mark.parametrize(
     "mask",
     [
@@ -318,9 +344,7 @@ def test_read_bounds_mask_invalid(naturalearth_lowres, mask):
         read_bounds(naturalearth_lowres, mask=mask)
 
 
-@pytest.mark.skipif(
-    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
-)
+@requires_shapely
 def test_read_bounds_bbox_mask_invalid(naturalearth_lowres):
     with pytest.raises(ValueError, match="cannot set both 'bbox' and 'mask'"):
         read_bounds(
@@ -328,9 +352,7 @@ def test_read_bounds_bbox_mask_invalid(naturalearth_lowres):
         )
 
 
-@pytest.mark.skipif(
-    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
-)
+@requires_shapely
 @pytest.mark.parametrize(
     "mask,expected",
     [
@@ -361,12 +383,8 @@ def test_read_bounds_mask(naturalearth_lowres_all_ext, mask, expected):
 
     fids = read_bounds(naturalearth_lowres_all_ext, mask=mask)[0]
 
-    if naturalearth_lowres_all_ext.suffix == ".gpkg":
-        # fid in gpkg is 1-based
-        assert array_equal(fids, np.array(expected) + 1)
-    else:
-        # fid in other formats is 0-based
-        assert array_equal(fids, expected)
+    fids_expected = np.array(expected) + START_FID[naturalearth_lowres_all_ext.suffix]
+    assert array_equal(fids, fids_expected)
 
 
 @pytest.mark.skipif(
@@ -382,21 +400,15 @@ def test_read_bounds_bbox_intersects_vs_envelope_overlaps(naturalearth_lowres_al
     if __gdal_geos_version__ is None:
         # bboxes for CAN, RUS overlap but do not intersect geometries
         assert fids.shape == (4,)
-        if naturalearth_lowres_all_ext.suffix == ".gpkg":
-            # fid in gpkg is 1-based
-            assert array_equal(fids, [4, 5, 19, 28])  # CAN, USA, RUS, MEX
-        else:
-            # fid in other formats is 0-based
-            assert array_equal(fids, [3, 4, 18, 27])  # CAN, USA, RUS, MEX
+        fids_expected = np.array([3, 4, 18, 27])  # CAN, USA, RUS, MEX
+        fids_expected += START_FID[naturalearth_lowres_all_ext.suffix]
+        assert array_equal(fids, fids_expected)
 
     else:
         assert fids.shape == (2,)
-        if naturalearth_lowres_all_ext.suffix == ".gpkg":
-            # fid in gpkg is 1-based
-            assert array_equal(fids, [5, 28])  # USA, MEX
-        else:
-            # fid in other formats is 0-based
-            assert array_equal(fids, [4, 27])  # USA, MEX
+        fids_expected = np.array([4, 27])  # USA, MEX
+        fids_expected += START_FID[naturalearth_lowres_all_ext.suffix]
+        assert array_equal(fids, fids_expected)
 
 
 @pytest.mark.parametrize("naturalearth_lowres", [".shp", ".gpkg"], indirect=True)
@@ -426,7 +438,8 @@ def test_read_info(naturalearth_lowres):
     elif naturalearth_lowres.suffix == ".shp":
         # fid_column == "" for formats where fid is not physically stored
         assert meta["fid_column"] == ""
-        # geometry_name == "" for formats where geometry column name cannot be customized
+        # geometry_name == "" for formats where geometry column name cannot be
+        # customized
         assert meta["geometry_name"] == ""
         assert meta["geometry_type"] == "Polygon"
         assert meta["driver"] == "ESRI Shapefile"
@@ -435,8 +448,14 @@ def test_read_info(naturalearth_lowres):
         raise ValueError(f"test not implemented for ext {naturalearth_lowres.suffix}")
 
 
-def test_read_info_vsi(naturalearth_lowres_vsi):
-    meta = read_info(naturalearth_lowres_vsi[1])
+@pytest.mark.parametrize(
+    "testfile", ["naturalearth_lowres_vsimem", "naturalearth_lowres_vsi"]
+)
+def test_read_info_vsi(testfile, request):
+    path = request.getfixturevalue(testfile)
+    path = path if not isinstance(path, tuple) else path[1]
+
+    meta = read_info(path)
 
     assert meta["fields"].shape == (5,)
     assert meta["features"] == 177
@@ -449,6 +468,13 @@ def test_read_info_bytes(geojson_bytes):
     assert meta["features"] == 3
 
 
+def test_read_info_nonseekable_bytes(nonseekable_bytes):
+    meta = read_info(nonseekable_bytes)
+
+    assert meta["fields"].shape == (0,)
+    assert meta["features"] == 1
+
+
 def test_read_info_filelike(geojson_filelike):
     meta = read_info(geojson_filelike)
 
@@ -483,8 +509,8 @@ def test_read_info_filelike(geojson_filelike):
         ),
     ],
 )
-def test_read_info_dataset_kwargs(data_dir, dataset_kwargs, fields):
-    meta = read_info(data_dir / "test_nested.geojson", **dataset_kwargs)
+def test_read_info_dataset_kwargs(nested_geojson_file, dataset_kwargs, fields):
+    meta = read_info(nested_geojson_file, **dataset_kwargs)
     assert meta["fields"].tolist() == fields
 
 
@@ -543,8 +569,8 @@ def test_read_info_unspecified_layer_warning(data_dir):
         read_info(data_dir / "sample.osm.pbf")
 
 
-def test_read_info_without_geometry(test_fgdb_vsi):
-    assert read_info(test_fgdb_vsi, layer="basetable_2")["total_bounds"] is None
+def test_read_info_without_geometry(no_geometry_file):
+    assert read_info(no_geometry_file)["total_bounds"] is None
 
 
 @pytest.mark.parametrize(
@@ -586,3 +612,67 @@ def test_error_handling_warning(capfd, naturalearth_lowres):
         read_info(naturalearth_lowres, INVALID="YES")
 
     assert capfd.readouterr().err == ""
+
+
+def test_vsimem_listtree_rmtree_unlink(naturalearth_lowres):
+    """Test all basic functionalities of file handling in /vsimem/."""
+    # Prepare test data in /vsimem
+    meta, _, geometry, field_data = read(naturalearth_lowres)
+    meta["spatial_index"] = False
+    meta["geometry_type"] = "MultiPolygon"
+    test_file_path = Path("/vsimem/pyogrio_test_naturalearth_lowres.gpkg")
+    test_dir_path = Path(f"/vsimem/pyogrio_dir_test/{naturalearth_lowres.stem}.gpkg")
+
+    write(test_file_path, geometry, field_data, **meta)
+    write(test_dir_path, geometry, field_data, **meta)
+
+    # Check if everything was created properly with listtree
+    files = vsi_listtree("/vsimem/")
+    assert test_file_path.as_posix() in files
+    assert test_dir_path.as_posix() in files
+
+    # Check listtree with pattern
+    files = vsi_listtree("/vsimem/", pattern="pyogrio_dir_test*.gpkg")
+    assert test_file_path.as_posix() not in files
+    assert test_dir_path.as_posix() in files
+
+    files = vsi_listtree("/vsimem/", pattern="pyogrio_test*.gpkg")
+    assert test_file_path.as_posix() in files
+    assert test_dir_path.as_posix() not in files
+
+    # Remove test_dir and its contents
+    vsi_rmtree(test_dir_path.parent)
+    files = vsi_listtree("/vsimem/")
+    assert test_file_path.as_posix() in files
+    assert test_dir_path.as_posix() not in files
+
+    # Remove test_file
+    vsi_unlink(test_file_path)
+
+
+def test_vsimem_rmtree_error(naturalearth_lowres_vsimem):
+    with pytest.raises(NotADirectoryError, match="Path is not a directory"):
+        vsi_rmtree(naturalearth_lowres_vsimem)
+
+    with pytest.raises(FileNotFoundError, match="Path does not exist"):
+        vsi_rmtree("/vsimem/non-existent")
+
+    with pytest.raises(
+        OSError, match="path to in-memory file or directory is required"
+    ):
+        vsi_rmtree("/vsimem")
+    with pytest.raises(
+        OSError, match="path to in-memory file or directory is required"
+    ):
+        vsi_rmtree("/vsimem/")
+
+    # Verify that naturalearth_lowres_vsimem still exists.
+    assert naturalearth_lowres_vsimem.as_posix() in vsi_listtree("/vsimem")
+
+
+def test_vsimem_unlink_error(naturalearth_lowres_vsimem):
+    with pytest.raises(IsADirectoryError, match="Path is a directory"):
+        vsi_unlink(naturalearth_lowres_vsimem.parent)
+
+    with pytest.raises(FileNotFoundError, match="Path does not exist"):
+        vsi_unlink("/vsimem/non-existent.gpkg")