[segyio] 121/376: Line._get_iter with range support for offsets

[Jørgen Kvalsvik]

This is an automated email from the git hooks/post-receive script.

jokva-guest pushed a commit to branch debian
in repository segyio.

commit 63dd89a016724fc391c7d148490a8121ec49bb96
Author: Jørgen Kvalsvik <jokva at statoil.com>
Date:   Mon Nov 14 10:51:44 2016 +0100

    Line._get_iter with range support for offsets
---
 python/segyio/_line.py | 48 +++++++++++++++++++++++++++++++++++++-----------
 tests/test_segy.py     | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/python/segyio/_line.py b/python/segyio/_line.py
index 53f9c97..b76ba69 100644
--- a/python/segyio/_line.py
+++ b/python/segyio/_line.py
@@ -55,23 +55,49 @@ class Line:
         t0 = self.trace0fn(lineno, offset)
         return self.readfn(t0, self.len, self.stride, buf)
 
+    # in order to support [:end] syntax, we must make sure
+    # start has a non-None value. lineno.indices() would set it
+    # to 0, but we don't know if that's a reasonable value or
+    # not. If start is None we set it to the first line
+    def _sanitize_slice(self, s, source):
+        if all((s.start, s.stop, s.step)):
+            return s
+
+        start, stop, step = s.start, s.stop, s.step
+        increasing = step is None or step > 0
+
+        if start is None:
+            start = source[0] if increasing else source[-1]
+
+        if stop is None:
+            stop = source[-1]+1 if increasing else source[0]-1
+
+        return slice(start, stop, step)
+
     def _get_iter(self, lineno, offset, buf):
         """ :rtype: collections.Iterable[numpy.ndarray]"""
 
-        # in order to support [:end] syntax, we must make sure
-        # start has a non-None value. lineno.indices() would set it
-        # to 0, but we don't know if that's a reasonable value or
-        # not. If start is None we set it to the first line
-        if lineno.start is None:
-            lineno = slice(self.lines[0], lineno.stop, lineno.step)
+        offsets = self.segy.offsets
+        if offset is None:
+            offset = offsets[0]
+
+        if not isinstance(lineno, slice):
+            lineno = slice(lineno, lineno + 1, 1)
+
+        if not isinstance(offset, slice):
+            offset = slice(offset, offset + 1, 1)
+
+        lineno = self._sanitize_slice(lineno, self.lines)
+        offset = self._sanitize_slice(offset, offsets)
 
         def gen():
-            o = self.segy.offsets[0]
-            s = set(self.lines)
-            rng = xrange(*lineno.indices(self.lines[-1] + 1))
+            offs, lns = set(self.segy.offsets), set(self.lines)
+            orng = xrange(*offset.indices(offsets[-1] + 1))
+            lrng = xrange(*lineno.indices(self.lines[-1] + 1))
 
-            for i in itertools.ifilter(s.__contains__, rng):
-                yield self._get(i, o, buf)
+            for l in itertools.ifilter(lns.__contains__, lrng):
+                for o in itertools.ifilter(offs.__contains__, orng):
+                    yield self._get(l, o, buf)
 
         return gen()
 
diff --git a/tests/test_segy.py b/tests/test_segy.py
index 6ed53ee..bfb7d0a 100644
--- a/tests/test_segy.py
+++ b/tests/test_segy.py
@@ -238,6 +238,51 @@ class TestSegy(TestCase):
             with self.assertRaises(TypeError):
                 f.iline[1, {}]
 
+    def test_iline_slice_fixed_offset(self):
+        with segyio.open(self.prestack, "r") as f:
+
+            for i, ln in enumerate(f.iline[:, 1], 1):
+                self.assertAlmostEqual(i + 100.01, ln[0][0], places = 4)
+                self.assertAlmostEqual(i + 100.02, ln[1][0], places = 4)
+                self.assertAlmostEqual(i + 100.03, ln[2][0], places = 4)
+
+                self.assertAlmostEqual(i + 100.01001, ln[0][1], places = 4)
+                self.assertAlmostEqual(i + 100.01002, ln[0][2], places = 4)
+                self.assertAlmostEqual(i + 100.02001, ln[1][1], places = 4)
+
+    def test_iline_slice_fixed_line(self):
+        with segyio.open(self.prestack, "r") as f:
+
+            for i, ln in enumerate(f.iline[1, :], 1):
+                off = i * 100
+                self.assertAlmostEqual(off + 1.01, ln[0][0], places = 4)
+                self.assertAlmostEqual(off + 1.02, ln[1][0], places = 4)
+                self.assertAlmostEqual(off + 1.03, ln[2][0], places = 4)
+
+                self.assertAlmostEqual(off + 1.01001, ln[0][1], places = 4)
+                self.assertAlmostEqual(off + 1.01002, ln[0][2], places = 4)
+                self.assertAlmostEqual(off + 1.02001, ln[1][1], places = 4)
+
+    def test_iline_slice_all_offsets(self):
+        with segyio.open(self.prestack, "r") as f:
+            offs, ils = len(f.offsets), len(f.ilines)
+            self.assertEqual(offs * ils, sum(1 for _ in f.iline[:,:]))
+            self.assertEqual(offs * ils, sum(1 for _ in f.iline[:,::-1]))
+            self.assertEqual(offs * ils, sum(1 for _ in f.iline[::-1,:]))
+            self.assertEqual(offs * ils, sum(1 for _ in f.iline[::-1,::-1]))
+            self.assertEqual(0,  sum(1 for _ in f.iline[:, 10:12]))
+            self.assertEqual(0,  sum(1 for _ in f.iline[10:12, :]))
+
+            self.assertEqual((offs / 2) * ils, sum(1 for _ in f.iline[::2, :]))
+            self.assertEqual(offs * (ils / 2), sum(1 for _ in f.iline[:, ::2]))
+
+            self.assertEqual((offs / 2) * ils, sum(1 for _ in f.iline[::-2, :]))
+            self.assertEqual(offs * (ils / 2), sum(1 for _ in f.iline[:, ::-2]))
+
+            self.assertEqual((offs / 2) * (ils / 2), sum(1 for _ in f.iline[::2, ::2]))
+            self.assertEqual((offs / 2) * (ils / 2), sum(1 for _ in f.iline[::2, ::-2]))
+            self.assertEqual((offs / 2) * (ils / 2), sum(1 for _ in f.iline[::-2, ::2]))
+
     def test_line_generators(self):
         with segyio.open(self.filename, "r") as f:
             for line in f.iline:

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/segyio.git