autotest: Add a VSIFile helper class

Author: dbaston

autotest/gcore/vsifile.py

@@ -37,7 +37,7 @@
 import pytest
 from lxml import etree
 
-from osgeo import gdal
+from osgeo import gdal, ogr
 
 
 ###############################################################################
@@ -1364,3 +1364,217 @@ def test_vsifile_CopyFileRestartable(tmp_vsimem):
     assert retcode == 0
     assert output_payload is None
     assert gdal.VSIStatL(dstfilename).size == 3
+
+
+###############################################################################
+# Test VSIFile helper class
+
+
+def test_vsifile_class_write_ascii(tmp_path):
+
+    fname = tmp_path / "test.txt"
+
+    lines = ["permission is hereby granted", "free of charge", "to any person"]
+
+    with gdaltest.vsi_open(fname, "w") as f:
+        assert f.tell() == 0
+
+        for line in lines:
+            f.write(line)
+            f.write("\n")
+
+    with open(fname, "r") as f:
+        assert [line.strip() for line in f.readlines()] == lines
+
+
+def test_vsifile_class_read_ascii(tmp_path):
+
+    fname = str(tmp_path / "test.txt")
+
+    lines = ["permission is hereby granted", "free of charge", "to any person"]
+
+    with open(fname, "w", newline="\n") as f:
+        for line in lines:
+            f.write(line)
+            f.write("\n")
+
+        with pytest.raises(Exception):
+            f.write(b"some bytes")
+
+    # read entire file
+    with gdaltest.vsi_open(fname, "r") as f:
+        contents = f.read()
+
+        assert type(contents) is str
+
+        lines_in = [line.strip() for line in contents.strip().split("\n")]
+        assert lines_in == lines
+
+    # read some characters
+    f = gdaltest.vsi_open(fname)
+    assert f.read(10) == "permission"
+
+    # skip a character
+    f.seek(1, os.SEEK_CUR)
+    assert f.read(9) == "is hereby"
+
+    f.seek(0, os.SEEK_SET)
+    assert f.read(10) == "permission"
+
+    # jump to end
+    f.seek(0, os.SEEK_END)
+    assert f.read(10) == ""
+
+    f.seek(-7, os.SEEK_END)
+    assert f.read() == "person\n"
+
+    f.close()
+    f.close()  # no harm in closing an already-closed file
+
+
+def test_vsifile_class_read_binary(tmp_path):
+
+    fname = tmp_path / "test.wkb"
+
+    g = ogr.CreateGeometryFromWkt("POINT (15 17)")
+    wkb = g.ExportToWkb()
+
+    with open(fname, "wb") as f:
+        f.write(wkb)
+
+    # read entire file
+    with gdaltest.vsi_open(fname, "rb") as f:
+        contents = f.read()
+
+        assert type(contents) is bytes
+
+        assert contents == wkb
+
+    # read some bytes
+    f = gdaltest.vsi_open(fname, "rb")
+    assert f.read(5) == wkb[:5]
+
+    f.seek(10, os.SEEK_SET)
+    assert f.read(5) == wkb[10:15]
+
+
+def test_vsifile_class_write_binary(tmp_path):
+
+    fname = tmp_path / "test.wkb"
+
+    g = ogr.CreateGeometryFromWkt("POINT (15 17)")
+    wkb = g.ExportToWkb()
+
+    with gdaltest.vsi_open(fname, "wb") as f:
+        f.write(wkb[:8])
+        f.write(wkb[8:])
+
+    with open(fname, "rb") as f:
+        assert f.read() == wkb
+
+
+def random_lines():
+    import random
+    import string
+
+    lines = []
+    for i in range(50):
+        lines.append(
+            "".join([random.choice(string.ascii_letters) for j in range(20 + 3 * i)])
+        )
+    lines.append(" ")
+    lines.append("")
+    lines.append("theend")
+    lines.append("")
+
+    return lines
+
+
+@pytest.mark.parametrize("terminating_newline", (True, False))
+def test_vsifile_class_line_iteration(tmp_path, terminating_newline):
+
+    fname = str(tmp_path / "test.txt")
+
+    lines_out = random_lines()
+
+    with open(fname, "w") as f:
+        for line in lines_out:
+            f.write(line)
+            f.write("\n")
+
+        if not terminating_newline:
+            f.write("lastline")
+            lines_out.append("lastline")
+
+    with gdaltest.vsi_open(fname) as f:
+        lines_in = [line for line in f]
+
+    assert lines_in == lines_out
+
+
+def test_vsifile_class_binary_line_iteration(tmp_path):
+
+    fname = str(tmp_path / "test.txt")
+
+    lines_out = [x.encode() for x in random_lines()]
+
+    with open(fname, "wb") as f:
+        for line in lines_out:
+            f.write(line)
+            f.write(b"\n")
+
+    with gdaltest.vsi_open(fname, "rb") as f:
+        lines_in = [line for line in f]
+
+    assert lines_in == lines_out
+
+
+def test_vsifile_class_zipped_csv_reader(tmp_path):
+
+    test_csv = str(tmp_path / "input.csv")
+    test_zip = str(tmp_path / "input.zip")
+
+    import csv
+    import shutil
+    import zipfile
+
+    shutil.copy("../ogr/data/prime_meridian.csv", test_csv)
+
+    with zipfile.ZipFile(test_zip, "w") as zf:
+        zf.write(test_csv, arcname="input.csv")
+
+    with gdaltest.vsi_open(f"/vsizip/{test_zip}/input.csv") as f:
+        records = [x for x in csv.DictReader(f)]
+
+    assert len(records) == 4
+    assert (
+        records[2]["INFORMATION_SOURCE"]
+        == "Institut Geographique National (IGN), Paris"
+    )
+
+
+def test_vsifile_class_file_does_not_exist(tmp_path):
+
+    with pytest.raises(OSError, match="No such file or directory"):
+        gdaltest.vsi_open(tmp_path / "does_not_exist.txt")
+
+
+def test_vsifile_class_read_from_closed_file(tmp_path):
+
+    with gdaltest.vsi_open(tmp_path / "out.txt", "w") as f:
+        f.write("abc")
+
+    with pytest.raises(ValueError, match="closed file"):
+        f.seek(0)
+
+
+def test_vsifile_class_append(tmp_vsimem):
+
+    fname = tmp_vsimem / "out.txt"
+
+    with gdaltest.vsi_open(fname, "w") as f:
+        f.write("abc")
+    with gdaltest.vsi_open(fname, "a") as f:
+        f.write("def")
+    with gdaltest.vsi_open(fname) as f:
+        assert f.read() == "abcdef"

autotest/pymod/gdaltest.py

@@ -2116,3 +2116,80 @@ def reopen(ds, update=False, open_options=None):
         allowed_drivers=[ds_drv.GetDescription()],
         open_options=open_options,
     )
+
+
+# VSIFile helper class
+
+
+class VSIFile:
+    def __init__(self, path, mode, encoding="utf-8"):
+        self._path = path
+        self._mode = mode
+
+        self._binary = "b" in mode
+        self._encoding = encoding
+
+        self._fp = gdal.VSIFOpenExL(self._path, self._mode, True)
+        if self._fp is None:
+            raise OSError(gdal.VSIGetLastErrorMsg())
+
+        self._closed = False
+        self._buffer = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        line = gdal.CPLReadLineL(self._fp)
+        if line is None:
+            raise StopIteration
+        if self._binary:
+            return line.encode()
+        return line
+
+    def close(self):
+        if self._closed:
+            return
+
+        self._closed = True
+        gdal.VSIFCloseL(self._fp)
+
+    def read(self, size=-1):
+        if size == -1:
+            pos = self.tell()
+            self.seek(0, 2)
+            size = self.tell()
+            self.seek(pos)
+
+        raw = gdal.VSIFReadL(1, size, self._fp)
+
+        if self._binary:
+            return bytes(raw)
+        else:
+            return raw.decode(self._encoding)
+
+    def write(self, x):
+
+        if self._binary:
+            assert type(x) in (bytes, bytearray, memoryview)
+        else:
+            assert type(x) is str
+            x = x.encode(self._encoding)
+
+        gdal.VSIFWriteL(x, 1, len(x), self._fp)
+
+    def seek(self, offset, whence=0):
+        gdal.VSIFSeekL(self._fp, offset, whence)
+
+    def tell(self):
+        return gdal.VSIFTellL(self._fp)
+
+
+def vsi_open(path, mode="r"):
+    return VSIFile(path, mode)

swig/include/cpl.i

@@ -943,6 +943,8 @@ int     VSIFWriteL( const char *, int, int, VSILFILE *fp );
 
 /* VSIFReadL() handled specially in python/gdal_python.i */
 
+const char* CPLReadLineL(VSILFILE* fp);
+
 void VSICurlClearCache();
 void VSICurlPartialClearCache( const char* utf8_path );