[segyio] 120/376: Line._get accepts offset argument

[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 1bcc199b7353c83335433045e86e3c1cf59a2859
Author: Jørgen Kvalsvik <jokva at statoil.com>
Date:   Thu Nov 10 17:29:32 2016 +0100

    Line._get accepts offset argument
---
 python/segyio/_line.py | 50 ++++++++++++++++++++++++++++++++++++--------------
 tests/test_segy.py     | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 14 deletions(-)

diff --git a/python/segyio/_line.py b/python/segyio/_line.py
index 7c02094..53f9c97 100644
--- a/python/segyio/_line.py
+++ b/python/segyio/_line.py
@@ -29,17 +29,33 @@ class Line:
         self.readfn = readfn
         self.writefn = writefn
 
-    def _get(self, lineno, buf):
+    def _get(self, lineno, offset, buf):
         """ :rtype: numpy.ndarray"""
+
+        offs = self.segy.offsets
+
+        if offset is None:
+            offset = 0
+        else:
+            try:
+                offset = next(i for i, x in enumerate(offs) if x == offset)
+            except StopIteration:
+                try:
+                    int(offset)
+                except TypeError:
+                    raise TypeError("Offset must be int or slice")
+
+                raise KeyError("Unknkown offset {}".format(offset))
+
         try:
             lineno = int(lineno)
         except TypeError:
             raise TypeError("Must be int or slice")
-        else:
-            t0 = self.trace0fn(lineno)
-            return self.readfn(t0, self.len, self.stride, buf)
 
-    def _get_iter(self, lineno, buf):
+        t0 = self.trace0fn(lineno, offset)
+        return self.readfn(t0, self.len, self.stride, buf)
+
+    def _get_iter(self, lineno, offset, buf):
         """ :rtype: collections.Iterable[numpy.ndarray]"""
 
         # in order to support [:end] syntax, we must make sure
@@ -50,25 +66,31 @@ class Line:
             lineno = slice(self.lines[0], lineno.stop, lineno.step)
 
         def gen():
+            o = self.segy.offsets[0]
             s = set(self.lines)
             rng = xrange(*lineno.indices(self.lines[-1] + 1))
 
             for i in itertools.ifilter(s.__contains__, rng):
-                yield self._get(i, buf)
+                yield self._get(i, o, buf)
 
         return gen()
 
-    def __getitem__(self, lineno):
+    def __getitem__(self, lineno, offset = None):
         """ :rtype: numpy.ndarray|collections.Iterable[numpy.ndarray]"""
         buf = self.buffn()
 
-        if isinstance(lineno, slice):
-            return self._get_iter(lineno, buf)
+        if isinstance(lineno, tuple):
+            lineno, offset = lineno
+
+        if isinstance(lineno, slice) or isinstance(offset, slice):
+            return self._get_iter(lineno, offset, buf)
 
-        return self._get(lineno, buf)
+        return self._get(lineno, offset, buf)
 
-    def trace0fn(self, lineno, offsets):
-        return segyio._segyio.fread_trace0(lineno, len(self.other_lines), self.stride, offsets, self.lines, self.name)
+    def trace0fn(self, lineno, offset_index):
+        offsets = len(self.segy.offsets)
+        trace0 = segyio._segyio.fread_trace0(lineno, len(self.other_lines), self.stride, offsets, self.lines, self.name)
+        return offset_index + trace0
 
     def __setitem__(self, lineno, val):
         if isinstance(lineno, slice):
@@ -83,7 +105,7 @@ class Line:
 
             return
 
-        t0 = self.trace0fn(lineno, len(self.segy.offsets))
+        t0 = self.trace0fn(lineno, 0)
         self.writefn(t0, self.len, self.stride, val)
 
     def __len__(self):
@@ -92,4 +114,4 @@ class Line:
     def __iter__(self):
         buf = self.buffn()
         for i in self.lines:
-            yield self._get(i, buf)
+            yield self._get(i, None, buf)
diff --git a/tests/test_segy.py b/tests/test_segy.py
index 729e834..6ed53ee 100644
--- a/tests/test_segy.py
+++ b/tests/test_segy.py
@@ -205,6 +205,39 @@ class TestSegy(TestCase):
             self.assertEqual(f.header[0][xl], f.header[1][xl])
             self.assertNotEqual(f.header[1][xl], f.header[2][xl])
 
+    def test_iline_offset(self):
+        with segyio.open(self.prestack, "r") as f:
+
+            line1 = f.iline[1, 1]
+            self.assertAlmostEqual(101.01, line1[0][0], places = 4)
+            self.assertAlmostEqual(101.02, line1[1][0], places = 4)
+            self.assertAlmostEqual(101.03, line1[2][0], places = 4)
+
+            self.assertAlmostEqual(101.01001, line1[0][1], places = 4)
+            self.assertAlmostEqual(101.01002, line1[0][2], places = 4)
+            self.assertAlmostEqual(101.02001, line1[1][1], places = 4)
+
+            line2 = f.iline[1, 2]
+            self.assertAlmostEqual(201.01, line2[0][0], places = 4)
+            self.assertAlmostEqual(201.02, line2[1][0], places = 4)
+            self.assertAlmostEqual(201.03, line2[2][0], places = 4)
+
+            self.assertAlmostEqual(201.01001, line2[0][1], places = 4)
+            self.assertAlmostEqual(201.01002, line2[0][2], places = 4)
+            self.assertAlmostEqual(201.02001, line2[1][1], places = 4)
+
+            with self.assertRaises(KeyError):
+                f.iline[1, 0]
+
+            with self.assertRaises(KeyError):
+                f.iline[1, 3]
+
+            with self.assertRaises(KeyError):
+                f.iline[100, 1]
+
+            with self.assertRaises(TypeError):
+                f.iline[1, {}]
+
     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