Merge pull request #100 from equinor/np.integer

Restore behaviour of accepting numpy integers as arguments to read functions

seismic_zfp/read.py

@@ -10,7 +10,7 @@
 from .version import SeismicZfpVersion
 from .utils import (pad, bytes_to_double, bytes_to_int, bytes_to_signed_int, get_chunk_cache_size,
                     coord_to_index, gen_coord_list, FileOffset, WrongDimensionalityError,
-                    get_correlated_diagonal_length, get_anticorrelated_diagonal_length)
+                    get_correlated_diagonal_length, get_anticorrelated_diagonal_length, python_int)
 import seismic_zfp
 from .sgzconstants import DISK_BLOCK_BYTES, SEGY_FILE_HEADER_BYTES, SEGY_TEXT_HEADER_BYTES
 from .headers import HeaderwordInfo
@@ -383,6 +383,7 @@ def read_inline(self, il_id):
         inline : numpy.ndarray of float32, shape: (n_xlines, n_samples)
             The specified inline, decompressed
         """
+        il_id = python_int(il_id)
         if self.is_2d:
             raise WrongDimensionalityError("Trying to read inlines from 2D file")
         if not 0 <= il_id < self.n_ilines:
@@ -426,6 +427,7 @@ def read_crossline(self, xl_id):
         crossline : numpy.ndarray of float32, shape: (n_ilines, n_samples)
             The specified crossline, decompressed
         """
+        xl_id = python_int(xl_id)
         if self.is_2d:
             raise WrongDimensionalityError("Trying to read crosslines from 2D file")
         if not 0 <= xl_id < self.n_xlines:
@@ -469,6 +471,7 @@ def read_zslice(self, zslice_id):
         zslice : numpy.ndarray of float32, shape: (n_ilines, n_xlines)
             The specified zslice (time or depth, depending on file contents), decompressed
         """
+        zslice_id = python_int(zslice_id)
         if self.is_2d:
             raise WrongDimensionalityError("Trying to read zslices from 2D file")
         if not 0 <= zslice_id < self.n_samples:
@@ -518,6 +521,7 @@ def read_correlated_diagonal(self, cd_id,
             - Shape (n_diagonal_traces OR max_cd_idx-min_cd_idx, n_samples OR max_sample_idx-min_sample_idx)
             The specified cd_slice, decompressed.
         """
+        cd_id = python_int(cd_id)
         if self.is_2d:
             raise WrongDimensionalityError("Trying to read diagonal from 2D file")
         if not -self.n_xlines < cd_id < self.n_ilines:
@@ -582,6 +586,7 @@ def read_anticorrelated_diagonal(self, ad_id,
             - Shape (n_diagonal_traces OR max_ad_idx-min_ad_idx, n_samples OR max_sample_idx-min_sample_idx)
             The specified ad_slice, decompressed.
         """
+        ad_id = python_int(ad_id)
         if self.is_2d:
             raise WrongDimensionalityError("Trying to read diagonal from 2D file")
         if not 0 <= ad_id < self.n_ilines + self.n_xlines - 1:
@@ -793,6 +798,7 @@ def get_trace(self, index, min_sample_id=None, max_sample_id=None, override_unst
         trace : numpy.ndarray of float32, shape (n_samples) or (max_sample_id - min_sample_id)
             A single trace, decompressed
         """
+        index = python_int(index)
         if self.is_2d:
             min_trace = self.blockshape[1] * (index // self.blockshape[1])
 

seismic_zfp/utils.py

@@ -150,6 +150,15 @@ def signed_int_to_bytes(bytes):
     return struct.pack('<i', bytes)
 
 
+def python_int(int_in):
+    if isinstance(int_in, int):
+        return int_in
+    elif isinstance(int_in, np.integer):
+        return int(int_in)
+    else:
+        raise TypeError("This function requires an integer argument")
+
+
 def define_blockshape_2d(bits_per_voxel, blockshape):
     assert blockshape[0] == 1
     return define_blockshape_3d(bits_per_voxel, blockshape)

tests/test_read.py

@@ -231,6 +231,23 @@ def test_read_inline():
     compare_inline_number(SGZ_FILE_8, SGY_FILE, [1, 2, 3, 4, 5], tolerance=1e-10)
 
 
+def test_numpy_ints():
+    sgz = seismic_zfp.open(SGZ_FILE_8)
+    for int_type in [np.int8, np.int16, np.uint8, np.uint16, np.intp, np.uintp]:
+        val = int_type(1)
+
+        sgz.read_inline(val)
+        sgz.read_crossline(val)
+        sgz.read_anticorrelated_diagonal(val)
+        sgz.read_correlated_diagonal(val)
+        sgz.get_trace(val)
+
+        sgz.iline[sgz.ilines[val]]
+        sgz.xline[sgz.xlines[val]]
+        sgz.depth_slice[val]
+        sgz.trace[val]
+
+
 def compare_crossline(sgz_filename, sgy_filename, lines, tolerance):
     with segyio.open(sgy_filename) as segyfile:
         for preload in [True, False]:

tests/test_utils.py

@@ -114,3 +114,11 @@ def test_define_blockshape_3d():
 def test_get_chunk_cache_size():
     assert 2048 == get_chunk_cache_size(1000, 2000)
     assert 1024 == get_chunk_cache_size(5000, 511)
+
+
+def test_python_int():
+    assert 1 == python_int(1)
+    assert 1 == python_int(np.uint16(1))
+    assert 1 == python_int(np.int8(1))
+    with pytest.raises(TypeError):
+        python_int("1")