[segyio] 279/376: Write range of traces; f.trace[:] = foo

[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 861e8d7392bd246f0521a064865f2e99d6188266
Author: Jørgen Kvalsvik <jokva at statoil.com>
Date:   Wed May 3 14:42:02 2017 +0200

    Write range of traces; f.trace[:] = foo
    
    Relax shape requirements on trace assignment, so that ranges of traces
    (or multi-dimensional arrays) can be written to file.
---
 python/segyio/_trace.py | 42 ++++++++++++++++++++----------------------
 python/test/segy.py     | 13 +++++++++++++
 2 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/python/segyio/_trace.py b/python/segyio/_trace.py
index 36de065..1f87665 100644
--- a/python/segyio/_trace.py
+++ b/python/segyio/_trace.py
@@ -1,6 +1,13 @@
 import numpy as np
 import segyio
 from segyio._raw_trace import RawTrace
+import itertools
+
+try:
+    from itertools import izip as zip
+    from itertools import imap as map
+except ImportError:  # will be 3.x series
+    pass
 
 
 class Trace:
@@ -37,25 +44,19 @@ class Trace:
         return self._readtr(start, 1, 1, buf)[0]
 
     def __setitem__(self, index, val):
-        if not 0 <= abs(index) < len(self):
-            raise IndexError("Trace %d not in range (-%d,%d)", (index, len(self), len(self)))
-
-        if not isinstance(val, np.ndarray):
-            raise TypeError("Value must be numpy.ndarray")
-
-        if val.dtype != np.single:
-            raise TypeError("Numpy array must be of type single")
-
-        shape = (len(self._file.samples),)
+        if isinstance(index, slice):
+            for i, x in zip(range(*index.indices(len(self))), val):
+                self.write_trace(i, x, self._file)
+            return
 
-        if val.shape[0] < shape[0]:
-            raise TypeError("Array wrong shape. Expected minimum %s, was %s" % (shape, val.shape))
+        if int(index) != index:
+            raise TypeError("Trace index must be integer type")
 
-        if not isinstance(index, slice):
-            index = slice(index, index + 1, 1)
+        if not 0 <= abs(index) < len(self):
+            raise IndexError("Trace %d not in range (-%d,%d)",
+                             (index, len(self), len(self)))
 
-        for i in range(*index.indices(len(self))):
-            self._writetr(i, val)
+        self.write_trace(index, val, self._file)
 
     def __len__(self):
         return self._file.tracecount
@@ -94,12 +95,6 @@ class Trace:
 
         return buf1, buf
 
-    def _writetr(self, traceno, buf):
-        if int(traceno) != traceno:
-            raise TypeError("Trace index must be integer type")
-
-        self.write_trace(int(traceno), buf, self._file)
-
     @classmethod
     def write_trace(cls, traceno, buf, segy):
         """
@@ -107,6 +102,9 @@ class Trace:
         :type buf: ?
         :type segy: segyio.SegyFile
         """
+        if isinstance(buf, np.ndarray) and buf.dtype != np.single:
+            raise TypeError("Numpy array must be of type single")
+
         segyio._segyio.write_trace(segy.xfd, traceno,
                                    buf,
                                    segy._tr0, segy._bsz,
diff --git a/python/test/segy.py b/python/test/segy.py
index 833e888..e270591 100644
--- a/python/test/segy.py
+++ b/python/test/segy.py
@@ -617,6 +617,19 @@ class TestSegy(TestCase):
             self.assertListEqual(list(f.attributes(189)[:]),
                                  [(i // 5) + 1 for i in range(len(f.trace))])
 
+    def test_assign_all_traces(self):
+        with TestContext("assign_all_traces") as context:
+            context.copy_file(self.filename, target = 'dst')
+            context.copy_file(self.filename)
+            with segyio.open('small.sgy') as f:
+                traces = f.trace.raw[:] * 2.0
+
+            with segyio.open('dst/small.sgy', 'r+') as f:
+                f.trace[:] = traces[:]
+
+            with segyio.open('dst/small.sgy') as f:
+                self.assertTrue(np.array_equal(f.trace.raw[:], traces))
+
     def test_traceaccess_from_array(self):
         a = np.arange(10, dtype = np.int)
         b = np.arange(10, dtype = np.int32)

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