pyogrio 0.9.0__cp310-cp310-macosx_12_0_arm64.whl

pyogrio/tests/test_core.py

@@ -0,0 +1,588 @@
+import numpy as np
+from numpy import array_equal, allclose
+import pytest
+
+from pyogrio import (
+    __gdal_version__,
+    __gdal_geos_version__,
+    list_drivers,
+    list_layers,
+    read_bounds,
+    read_info,
+    set_gdal_config_options,
+    get_gdal_config_option,
+    get_gdal_data_path,
+)
+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
+
+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
+
+
+try:
+    import shapely
+except ImportError:
+    pass
+
+
+def test_gdal_data():
+    # test will fail if GDAL data files cannot be found, indicating an
+    # installation error
+    assert has_gdal_data()
+
+
+def test_proj_data():
+    # test will fail if PROJ data files cannot be found, indicating an
+    # installation error
+    assert has_proj_data()
+
+
+def test_get_gdal_data_path():
+    # test will fail if the function returns None, which means that GDAL
+    # cannot find data files, indicating an installation error
+    assert isinstance(get_gdal_data_path(), str)
+
+
+def test_gdal_geos_version():
+    assert __gdal_geos_version__ is None or isinstance(__gdal_geos_version__, tuple)
+
+
+@pytest.mark.parametrize(
+    "path,expected",
+    [
+        ("test.shp", "ESRI Shapefile"),
+        ("test.shp.zip", "ESRI Shapefile"),
+        ("test.geojson", "GeoJSON"),
+        ("test.geojsonl", "GeoJSONSeq"),
+        ("test.gpkg", "GPKG"),
+        pytest.param(
+            "test.gpkg.zip",
+            "GPKG",
+            marks=pytest.mark.skipif(
+                __gdal_version__ < (3, 7, 0),
+                reason="writing *.gpkg.zip requires GDAL >= 3.7.0",
+            ),
+        ),
+        # postgres can be detected by prefix instead of extension
+        pytest.param(
+            "PG:dbname=test",
+            "PostgreSQL",
+            marks=pytest.mark.skipif(
+                "PostgreSQL" not in list_drivers(),
+                reason="PostgreSQL path test requires PostgreSQL driver",
+            ),
+        ),
+    ],
+)
+def test_detect_write_driver(path, expected):
+    assert detect_write_driver(path) == expected
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "test.svg",  # only supports read
+        "test.",  # not a valid extension
+        "test",  # no extension or prefix
+        "test.foo",  # not a valid extension
+        "FOO:test",  # not a valid prefix
+    ],
+)
+def test_detect_write_driver_unsupported(path):
+    with pytest.raises(ValueError, match="Could not infer driver from path"):
+        detect_write_driver(path)
+
+
+@pytest.mark.parametrize("path", ["test.xml", "test.txt"])
+def test_detect_write_driver_multiple_unsupported(path):
+    with pytest.raises(ValueError, match="multiple drivers are available"):
+        detect_write_driver(path)
+
+
+@pytest.mark.parametrize(
+    "driver,expected",
+    [
+        # drivers known to be well-supported by pyogrio
+        ("ESRI Shapefile", True),
+        ("GeoJSON", True),
+        ("GeoJSONSeq", True),
+        ("GPKG", True),
+        # drivers not supported for write by GDAL
+        ("HTTP", False),
+        ("OAPIF", False),
+    ],
+)
+def test_ogr_driver_supports_write(driver, expected):
+    assert ogr_driver_supports_write(driver) == expected
+
+
+def test_list_drivers():
+    all_drivers = list_drivers()
+
+    # verify that the core drivers are present
+    for name in ("ESRI Shapefile", "GeoJSON", "GeoJSONSeq", "GPKG", "OpenFileGDB"):
+        assert name in all_drivers
+
+        expected_capability = "rw"
+        if name == "OpenFileGDB" and __gdal_version__ < (3, 6, 0):
+            expected_capability = "r"
+
+        assert all_drivers[name] == expected_capability
+
+    drivers = list_drivers(read=True)
+    expected = {k: v for k, v in all_drivers.items() if v.startswith("r")}
+    assert len(drivers) == len(expected)
+
+    drivers = list_drivers(write=True)
+    expected = {k: v for k, v in all_drivers.items() if v.endswith("w")}
+    assert len(drivers) == len(expected)
+
+    drivers = list_drivers(read=True, write=True)
+    expected = {
+        k: v for k, v in all_drivers.items() if v.startswith("r") and v.endswith("w")
+    }
+    assert len(drivers) == len(expected)
+
+
+def test_list_layers(naturalearth_lowres, naturalearth_lowres_vsi, test_fgdb_vsi):
+    assert array_equal(
+        list_layers(naturalearth_lowres), [["naturalearth_lowres", "Polygon"]]
+    )
+
+    assert array_equal(
+        list_layers(naturalearth_lowres_vsi[1]), [["naturalearth_lowres", "Polygon"]]
+    )
+
+    # 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
+
+        # Make sure that nonspatial layer has None for geometry
+        assert array_equal(fgdb_layers[0], ["basetable_2", None])
+
+        # 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"])
+
+
+def test_list_layers_bytes(geojson_bytes):
+    layers = list_layers(geojson_bytes)
+
+    assert layers.shape == (1, 2)
+    assert layers[0, 0] == "test"
+
+
+def test_list_layers_filelike(geojson_filelike):
+    layers = list_layers(geojson_filelike)
+
+    assert layers.shape == (1, 2)
+    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])
+
+
+def test_read_bounds_vsi(naturalearth_lowres_vsi):
+    fids, bounds = read_bounds(naturalearth_lowres_vsi[1])
+    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])
+
+
+def test_read_bounds_bytes(geojson_bytes):
+    fids, bounds = read_bounds(geojson_bytes)
+    assert fids.shape == (3,)
+    assert bounds.shape == (4, 3)
+    assert allclose(bounds[:, 0], [-180.0, -18.28799, 180.0, -16.02088])
+
+
+def test_read_bounds_filelike(geojson_filelike):
+    fids, bounds = read_bounds(geojson_filelike)
+    assert fids.shape == (3,)
+    assert bounds.shape == (4, 3)
+    assert allclose(bounds[:, 0], [-180.0, -18.28799, 180.0, -16.02088])
+
+
+def test_read_bounds_max_features(naturalearth_lowres):
+    bounds = read_bounds(naturalearth_lowres, max_features=2)[1]
+    assert bounds.shape == (4, 2)
+
+
+def test_read_bounds_unspecified_layer_warning(data_dir):
+    """Reading a multi-layer file without specifying a layer gives a warning."""
+    with pytest.warns(UserWarning, match="More than one layer found "):
+        read_bounds(data_dir / "sample.osm.pbf")
+
+
+def test_read_bounds_negative_max_features(naturalearth_lowres):
+    with pytest.raises(ValueError, match="'max_features' must be >= 0"):
+        read_bounds(naturalearth_lowres, max_features=-1)
+
+
+def test_read_bounds_skip_features(naturalearth_lowres):
+    expected_bounds = read_bounds(naturalearth_lowres, max_features=11)[1][:, 10]
+    fids, bounds = read_bounds(naturalearth_lowres, skip_features=10)
+    assert bounds.shape == (4, 167)
+    assert allclose(bounds[:, 0], expected_bounds)
+    assert fids[0] == 10
+
+
+def test_read_bounds_negative_skip_features(naturalearth_lowres):
+    with pytest.raises(ValueError, match="'skip_features' must be >= 0"):
+        read_bounds(naturalearth_lowres, skip_features=-1)
+
+
+def test_read_bounds_where_invalid(naturalearth_lowres_all_ext):
+    with pytest.raises(ValueError, match="Invalid SQL"):
+        read_bounds(naturalearth_lowres_all_ext, where="invalid")
+
+
+def test_read_bounds_where(naturalearth_lowres):
+    fids, bounds = read_bounds(naturalearth_lowres, where="iso_a3 = 'CAN'")
+    assert fids.shape == (1,)
+    assert bounds.shape == (4, 1)
+    assert fids[0] == 3
+    assert allclose(bounds[:, 0], [-140.99778, 41.675105, -52.648099, 83.23324])
+
+
+@pytest.mark.parametrize("bbox", [(1,), (1, 2), (1, 2, 3)])
+def test_read_bounds_bbox_invalid(naturalearth_lowres, bbox):
+    with pytest.raises(ValueError, match="Invalid bbox"):
+        read_bounds(naturalearth_lowres, bbox=bbox)
+
+
+def test_read_bounds_bbox(naturalearth_lowres_all_ext):
+    # should return no features
+    fids, bounds = read_bounds(
+        naturalearth_lowres_all_ext, bbox=(0, 0, 0.00001, 0.00001)
+    )
+
+    assert fids.shape == (0,)
+    assert bounds.shape == (4, 0)
+
+    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
+
+    assert bounds.shape == (4, 2)
+    assert allclose(
+        bounds.T,
+        [
+            [-82.96578305, 7.22054149, -77.24256649, 9.61161001],
+            [-85.94172543, 8.22502798, -82.54619626, 11.21711925],
+        ],
+    )
+
+
+@pytest.mark.skipif(
+    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
+)
+@pytest.mark.parametrize(
+    "mask",
+    [
+        {"type": "Point", "coordinates": [0, 0]},
+        '{"type": "Point", "coordinates": [0, 0]}',
+        "invalid",
+    ],
+)
+def test_read_bounds_mask_invalid(naturalearth_lowres, mask):
+    with pytest.raises(ValueError, match="'mask' parameter must be a Shapely geometry"):
+        read_bounds(naturalearth_lowres, mask=mask)
+
+
+@pytest.mark.skipif(
+    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
+)
+def test_read_bounds_bbox_mask_invalid(naturalearth_lowres):
+    with pytest.raises(ValueError, match="cannot set both 'bbox' and 'mask'"):
+        read_bounds(
+            naturalearth_lowres, bbox=(-85, 8, -80, 10), mask=shapely.Point(-105, 55)
+        )
+
+
+@pytest.mark.skipif(
+    not HAS_SHAPELY, reason="Shapely is required for mask functionality"
+)
+@pytest.mark.parametrize(
+    "mask,expected",
+    [
+        ("POINT (-105 55)", [3]),
+        ("POLYGON ((-80 8, -80 10, -85 10, -85 8, -80 8))", [33, 34]),
+        (
+            """POLYGON ((
+                6.101929 50.97085,
+                5.773002 50.906611,
+                5.593156 50.642649,
+                6.059271 50.686052,
+                6.374064 50.851481,
+                6.101929 50.97085
+            ))""",
+            [121, 129, 130],
+        ),
+        (
+            """GEOMETRYCOLLECTION (
+                POINT (-7.7 53),
+                POLYGON ((-80 8, -80 10, -85 10, -85 8, -80 8))
+            )""",
+            [33, 34, 133],
+        ),
+    ],
+)
+def test_read_bounds_mask(naturalearth_lowres_all_ext, mask, expected):
+    mask = shapely.from_wkt(mask)
+
+    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)
+
+
+@pytest.mark.skipif(
+    __gdal_version__ < (3, 4, 0),
+    reason="Cannot determine if GEOS is present or absent for GDAL < 3.4",
+)
+def test_read_bounds_bbox_intersects_vs_envelope_overlaps(naturalearth_lowres_all_ext):
+    # If GEOS is present and used by GDAL, bbox filter will be based on intersection
+    # of bbox and actual geometries; if GEOS is absent or not used by GDAL, it
+    # will be based on overlap of bounding boxes instead
+    fids, _ = read_bounds(naturalearth_lowres_all_ext, bbox=(-140, 20, -100, 45))
+
+    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
+
+    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
+
+
+@pytest.mark.parametrize("naturalearth_lowres", [".shp", ".gpkg"], indirect=True)
+def test_read_info(naturalearth_lowres):
+    meta = read_info(naturalearth_lowres)
+
+    assert meta["layer_name"] == "naturalearth_lowres"
+    assert meta["crs"] == "EPSG:4326"
+    assert meta["encoding"] == "UTF-8"
+    assert meta["fields"].shape == (5,)
+    assert meta["dtypes"].tolist() == ["int64", "object", "object", "object", "float64"]
+    assert meta["features"] == 177
+    assert allclose(meta["total_bounds"], (-180, -90, 180, 83.64513))
+    assert meta["capabilities"]["random_read"] is True
+    assert meta["capabilities"]["fast_spatial_filter"] is False
+    assert meta["capabilities"]["fast_feature_count"] is True
+    assert meta["capabilities"]["fast_total_bounds"] is True
+
+    if naturalearth_lowres.suffix == ".gpkg":
+        assert meta["fid_column"] == "fid"
+        assert meta["geometry_name"] == "geom"
+        assert meta["geometry_type"] == "MultiPolygon"
+        assert meta["driver"] == "GPKG"
+        if GDAL_GE_38:
+            # this capability is only True for GPKG if GDAL >= 3.8
+            assert meta["capabilities"]["fast_set_next_by_index"] is True
+    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
+        assert meta["geometry_name"] == ""
+        assert meta["geometry_type"] == "Polygon"
+        assert meta["driver"] == "ESRI Shapefile"
+        assert meta["capabilities"]["fast_set_next_by_index"] is True
+    else:
+        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])
+
+    assert meta["fields"].shape == (5,)
+    assert meta["features"] == 177
+
+
+def test_read_info_bytes(geojson_bytes):
+    meta = read_info(geojson_bytes)
+
+    assert meta["fields"].shape == (5,)
+    assert meta["features"] == 3
+
+
+def test_read_info_filelike(geojson_filelike):
+    meta = read_info(geojson_filelike)
+
+    assert meta["fields"].shape == (5,)
+    assert meta["features"] == 3
+
+
+@pytest.mark.parametrize(
+    "dataset_kwargs,fields",
+    [
+        ({}, ["top_level", "intermediate_level"]),
+        (
+            {"FLATTEN_NESTED_ATTRIBUTES": "YES"},
+            [
+                "top_level",
+                "intermediate_level_bottom_level",
+            ],
+        ),
+        (
+            {"flatten_nested_attributes": "yes"},
+            [
+                "top_level",
+                "intermediate_level_bottom_level",
+            ],
+        ),
+        (
+            {"flatten_nested_attributes": True},
+            [
+                "top_level",
+                "intermediate_level_bottom_level",
+            ],
+        ),
+    ],
+)
+def test_read_info_dataset_kwargs(data_dir, dataset_kwargs, fields):
+    meta = read_info(data_dir / "test_nested.geojson", **dataset_kwargs)
+    assert meta["fields"].tolist() == fields
+
+
+def test_read_info_invalid_dataset_kwargs(naturalearth_lowres):
+    with pytest.warns(RuntimeWarning, match="does not support open option INVALID"):
+        read_info(naturalearth_lowres, INVALID="YES")
+
+
+def test_read_info_force_feature_count_exception(data_dir):
+    with pytest.raises(DataLayerError, match="Could not iterate over features"):
+        read_info(data_dir / "sample.osm.pbf", layer="lines", force_feature_count=True)
+
+
+@pytest.mark.parametrize(
+    "layer, force, expected",
+    [
+        ("points", False, -1),
+        ("points", True, 8),
+        ("lines", False, -1),
+        ("lines", True, 36),
+    ],
+)
+def test_read_info_force_feature_count(data_dir, layer, force, expected):
+    # the sample OSM file has non-increasing node IDs which causes the default
+    # custom indexing to raise an exception iterating over features
+    meta = read_info(
+        data_dir / "sample.osm.pbf",
+        layer=layer,
+        force_feature_count=force,
+        USE_CUSTOM_INDEXING=False,
+    )
+    assert meta["features"] == expected
+
+
+@pytest.mark.parametrize(
+    "force_total_bounds, expected_total_bounds",
+    [(True, (-180.0, -90.0, 180.0, 83.64513)), (False, None)],
+)
+def test_read_info_force_total_bounds(
+    tmp_path, naturalearth_lowres, force_total_bounds, expected_total_bounds
+):
+    geojson_path = prepare_testfile(
+        naturalearth_lowres, dst_dir=tmp_path, ext=".geojsonl"
+    )
+
+    info = read_info(geojson_path, force_total_bounds=force_total_bounds)
+    if expected_total_bounds is not None:
+        assert allclose(info["total_bounds"], expected_total_bounds)
+    else:
+        assert info["total_bounds"] is None
+
+
+def test_read_info_unspecified_layer_warning(data_dir):
+    """Reading a multi-layer file without specifying a layer gives a warning."""
+    with pytest.warns(UserWarning, match="More than one layer found "):
+        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
+
+
+@pytest.mark.parametrize(
+    "name,value,expected",
+    [
+        ("CPL_DEBUG", "ON", True),
+        ("CPL_DEBUG", True, True),
+        ("CPL_DEBUG", "OFF", False),
+        ("CPL_DEBUG", False, False),
+    ],
+)
+def test_set_config_options(name, value, expected):
+    set_gdal_config_options({name: value})
+    actual = get_gdal_config_option(name)
+    assert actual == expected
+
+
+def test_reset_config_options():
+    set_gdal_config_options({"foo": "bar"})
+    assert get_gdal_config_option("foo") == "bar"
+
+    set_gdal_config_options({"foo": None})
+    assert get_gdal_config_option("foo") is None
+
+
+def test_error_handling(capfd):
+    # an operation that triggers a GDAL Failure
+    # -> error translated into Python exception + not printed to stderr
+    with pytest.raises(DataSourceError, match="No such file or directory"):
+        read_info("non-existent.shp")
+
+    assert capfd.readouterr().err == ""
+
+
+def test_error_handling_warning(capfd, naturalearth_lowres):
+    # an operation that triggers a GDAL Warning
+    # -> translated into a Python warning + not printed to stderr
+    with pytest.warns(RuntimeWarning, match="does not support open option INVALID"):
+        read_info(naturalearth_lowres, INVALID="YES")
+
+    assert capfd.readouterr().err == ""