Merge pull request #632 from achaikou/python_datasource

Expose datasource to python

Author: achaikou

.github/workflows/build.yaml

@@ -134,3 +134,31 @@ jobs:
         uses: matlab-actions/run-command@v2
         with:
           command: addpath('build/mex'), try, run('build/mex/test/segyspec.m'), exit(0), catch ME, disp(getReport(ME, 'extended')), exit(-1), end;
+
+  python:
+    runs-on: ubuntu-latest
+    name: Build and test python on example stream
+    strategy:
+      fail-fast: false
+      matrix:
+        stream: [pyfile, pymemory]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install build/test dependencies
+        working-directory: python
+        run: |
+          python3 -m pip install -r requirements-dev.txt
+
+      - name: Configure and build
+        shell: bash
+        run: |
+          cmake -S . -B build -DBUILD_PYTHON=ON -DBUILD_BIN=OFF
+          cmake --build build
+
+      - name: Test
+        shell: bash
+        run: |
+          cd python
+          pytest -rP --datasource ${{ matrix.stream }} test/stream.py

python/segyio/__init__.py

@@ -88,8 +88,8 @@ def enums(cls):
 from .tracesortingformat import TraceSortingFormat
 from .tracefield import TraceField
 from . import su
-from .open import open
-from .create import create
+from .open import open, open_with, open_from_memory
+from .create import create, create_with
 from .segy import SegyFile, spec
 from .tools import dt, sample_indexes, create_text_header, native
 from .tools import collect, cube

python/segyio/create.py

@@ -3,7 +3,10 @@
 import segyio
 
 from . import TraceSortingFormat
-from .utils import c_endianness
+from .utils import (
+    FileDatasourceDescriptor,
+    StreamDatasourceDescriptor,
+)
 
 def default_text_header(iline, xline, offset):
     lines = {
@@ -178,9 +181,42 @@ def create(filename, spec):
     ...         dst.header = src.header
     ...         dst.trace = src.trace
     """
+    return _create(FileDatasourceDescriptor(filename, "w+"), spec)
 
-    from . import _segyio
 
+def create_with(stream, spec, minimize_requests_number=True):
+    """
+    Creates a segy file on stream.
+
+    Function behaves the same as `segyio.create`, but outputs data to a
+    finite stream instead of a file. Stream's close() will be called when
+    SegyFile is closed.
+
+    Note that `segyio.create_with` can be very slow. `segyio.create` is
+    generally a preferred option when speed matters.
+
+    Parameters
+    ----------
+
+    stream : file-like object
+        Data destination. It is up to the user to assure stream is opened in w+b mode.
+    spec : segyio.spec
+        Structure of the segy file
+    minimize_requests_number : bool
+        Configuration for some internal algorithms. True to minimize number of
+        requests to the stream at the cost of higher memory usage. False to
+        minimize memory usage at the cost of more requests to the stream.
+    """
+    return _create(
+        StreamDatasourceDescriptor(
+            stream,
+            minimize_requests_number
+        ),
+        spec
+    )
+
+
+def _create(datasource_descriptor, spec):
     if not structured(spec):
         tracecount = spec.tracecount
     else:
@@ -193,9 +229,7 @@ def create(filename, spec):
     if endian is None:
         endian = 'big'
 
-    fd = _segyio.segyfd(
-        filename=str(filename), mode='w+', endianness=c_endianness(endian)
-    )
+    fd = datasource_descriptor.make_segyfile_descriptor(endian)
     fd.segymake(
         samples = len(samples),
         tracecount = tracecount,
@@ -204,8 +238,7 @@ def create(filename, spec):
     )
 
     f = segyio.SegyFile(fd,
-            filename = str(filename),
-            mode = 'w+',
+            datasource_descriptor,
             iline = int(spec.iline),
             xline = int(spec.xline),
             endian = endian,

python/segyio/open.py

@@ -2,7 +2,11 @@
 
 import segyio
 
-from .utils import c_endianness
+from .utils import (
+    FileDatasourceDescriptor,
+    StreamDatasourceDescriptor,
+    MemoryBufferDatasourceDescriptor
+)
 
 def infer_geometry(f, metrics, iline, xline, strict):
     try:
@@ -148,16 +152,95 @@ def open(filename, mode="r", iline = 189,
         solution = 'use r+ to open in read-write'
         raise ValueError(', '.join((problem, solution)))
 
-    from . import _segyio
-    fd = _segyio.segyfd(
-        filename=str(filename), mode=mode, endianness=c_endianness(endian)
+    return _open(
+        FileDatasourceDescriptor(filename, mode),
+        iline, xline, strict, ignore_geometry, endian
     )
+
+
+def open_with(stream,
+              iline=189,
+              xline=193,
+              strict=True,
+              ignore_geometry=False,
+              endian='big',
+              minimize_requests_number=True
+              ):
+    """
+    Opens a segy file from a stream.
+
+    Function behaves the same as `segyio.open`, but sources data from a finite
+    stream instead of a file. Stream's close() will be called when SegyFile is
+    closed.
+
+    Note that `segyio.open_with` can be very slow. `segyio.open` is generally a
+    preferred option when speed matters.
+
+    Parameters
+    ----------
+
+    stream : file-like object
+        Source of the data. It is up to the user to assure stream is opened in rb
+        mode for reading and r+b mode for updating.
+    minimize_requests_number : bool
+        Configuration for some internal algorithms. True to minimize number of
+        requests to the stream at the cost of higher memory usage. False to
+        minimize memory usage at the cost of more requests to the stream.
+
+    See other common parameters at `segyio.open`.
+    """
+    return _open(
+        StreamDatasourceDescriptor(
+            stream,
+            minimize_requests_number
+        ),
+        iline, xline, strict, ignore_geometry, endian
+    )
+
+
+def open_from_memory(memory_buffer,
+                   iline=189,
+                   xline=193,
+                   strict=True,
+                   ignore_geometry=False,
+                   endian='big'):
+    """
+    Opens a segy file from memory.
+
+    Function behaves the same as `segyio.open`, but expects file to fit
+    completely into memory. Performance is expected to be significantly faster
+    than using `segyio.open` or `segyio.open_with`.
+
+    Parameters
+    ----------
+
+    buffer : contiguous memory object
+        Source of the data: complete SEGY file loaded into memory. Update
+        operations assume memory is writable.
+
+    See other common parameters at `segyio.open`.
+    """
+    return _open(
+        MemoryBufferDatasourceDescriptor(
+            memory_buffer
+        ),
+        iline, xline, strict, ignore_geometry, endian
+    )
+
+
+def _open(datasource_descriptor,
+          iline=189,
+          xline=193,
+          strict=True,
+          ignore_geometry=False,
+          endian='big'):
+
+    fd = datasource_descriptor.make_segyfile_descriptor(endian)
     fd.segyopen()
     metrics = fd.metrics()
 
     f = segyio.SegyFile(fd,
-            filename = str(filename),
-            mode = mode,
+            datasource_descriptor,
             iline = int(iline),
             xline = int(xline),
             endian = endian,

python/segyio/segy.py

@@ -23,13 +23,16 @@ class SegyFile(object):
 
     _unstructured_errmsg = "File opened in unstructured mode."
 
-    def __init__(self, fd, filename, mode, iline=189,
-                                           xline=193,
-                                           endian='big',
-                                           ):
-
-        self._filename = filename
-        self._mode = mode
+    def __init__(
+        self,
+        fd,
+        datasource_descriptor,
+        iline=189,
+        xline=193,
+        endian='big'
+    ):
+
+        self._datasource_descriptor = datasource_descriptor
         self._il = iline
         self._xl = xline
 
@@ -90,7 +93,7 @@ def __init__(self, fd, filename, mode, iline=189,
         super(SegyFile, self).__init__()
 
     def __str__(self):
-        f = "SegyFile {}:".format(self._filename)
+        f = "SegyFile {}:".format(str(self._datasource_descriptor))
 
         if self.unstructured:
             il =  "  inlines: None"
@@ -115,8 +118,8 @@ def __str__(self):
 
 
     def __repr__(self):
-        return "SegyFile('{}', '{}', iline = {}, xline = {})".format(
-                        self._filename, self._mode, self._il, self._xl)
+        return "SegyFile('{}', iline = {}, xline = {})".format(
+                        self._datasource_descriptor.__repr__(), self._il, self._xl)
 
     def __enter__(self):
         return self
@@ -853,7 +856,7 @@ def readonly(self):
         .. versionadded:: 1.6
         """
 
-        return '+' not in self._mode
+        return self._datasource_descriptor.readonly()
 
 
     def interpret(self, ilines, xlines, offsets=None, sorting=TraceSortingFormat.INLINE_SORTING):

python/segyio/segyio.cpp

@@ -466,6 +466,10 @@ int init( segyfd* self, PyObject* args, PyObject* kwargs ) {
 
     autods ds;
     if( stream ) {
+        if (stream == Py_None) {
+            ValueError( "stream must not be None" );
+            return -1;
+        }
         if( minimize_requests_number < 0 ) {
             ValueError( "minimize requests number is not set" );
             return -1;
@@ -474,6 +478,10 @@ int init( segyfd* self, PyObject* args, PyObject* kwargs ) {
             stream, minimize_requests_number
         );
     } else if( memory_buffer ) {
+        if (memory_buffer == Py_None) {
+            ValueError( "buffer must not be None" );
+            return -1;
+        }
         auto flags = PyBUF_CONTIG | PyBUF_WRITABLE;
         const int err = PyObject_GetBuffer( memory_buffer, &ds.memory_ds_buffer, flags );
         if( err ) {

python/segyio/su/file.py

@@ -1,7 +1,7 @@
 from ..open import infer_geometry
 from ..segy import SegyFile
 from . import words
-from ..utils import c_endianness
+from ..utils import FileDatasourceDescriptor
 
 import numpy
 
@@ -81,17 +81,26 @@ def open(filename, mode = 'r', iline = 189,
         solution = 'use r+ to open in read-write'
         raise ValueError(', '.join((problem, solution)))
 
-    from .. import _segyio
-    fd = _segyio.segyfd(
-        filename=str(filename), mode=mode, endianness=c_endianness(endian)
+    return _open(
+        FileDatasourceDescriptor(filename, mode),
+        iline, xline, strict, ignore_geometry, endian
     )
+
+
+def _open(datasource_descriptor,
+          iline=189,
+          xline=193,
+          strict=True,
+          ignore_geometry=False,
+          endian='big'):
+
+    fd = datasource_descriptor.make_segyfile_descriptor(endian)
     fd.suopen()
     metrics = fd.metrics()
 
     f = sufile(
         fd,
-        filename = str(filename),
-        mode = mode,
+        datasource_descriptor,
         iline = iline,
         xline = xline,
     )