Support attributes operations on tracefield

Author: achaikou

python/docs/segyio.rst

@@ -96,6 +96,7 @@ Trace and binary header
     :exclude-members: count
 
 .. autoclass:: segyio.trace.FieldLayoutEntry()
+    :special-members: __getitem__
 
 Tools
 =====

python/examples/flip-endianness.py

@@ -32,20 +32,12 @@ def main():
             dst.bin = src.bin
             dst.trace = src.trace
 
-            # intention:
-            # dst.traceheader = src.traceheader
-
         # traceheader mapping was copied after open. At the moment of creation
         # of this example file code can't deal with that, so we must reopen the
         # file after bin and ext_headers are copied so that mapping is
         # recognized.
         with segyio.open(dstfile, "r+", endian="little", ignore_geometry=True) as dst:
-            for trace_index in range(src.tracecount):
-                for traceheader_index in range(src.traceheader_count):
-                    traceheader = src.segyfd.getth(
-                        trace_index, traceheader_index, bytearray(240))
-                    dst.segyfd.putth(
-                        trace_index, traceheader_index, traceheader)
+            dst.traceheader = src.traceheader
 
 
 if __name__ == '__main__':

python/segyio/segy.py

@@ -414,6 +414,16 @@ def tracefield(self):
         >>> f.tracefield.SEG00000.iline.offset()
         >>> f.tracefield[0].xline.offset()
 
+        Get values of field for many traces:
+
+        >>> f.tracefield.SEG00001.dt[0, 10, 20]
+
+        Note:
+
+        - this works similar to :meth:`.attributes`
+        - use :meth:`.traceheader` if you need to get many fields from a single
+          trace.
+
         Notes
         -----
         .. versionadded:: 2.0
@@ -546,8 +556,7 @@ def attributes(self, field):
         .. versionadded:: 1.1
 
         """
-        traceheader_layout = self._traceheader_layouts["SEG00000"]
-        return Attributes(field, self.segyfd, traceheader_layout, self.tracecount)
+        return Attributes(self, field, 0)
 
     @property
     def trace(self):

python/segyio/trace.py

@@ -1057,12 +1057,14 @@ class Attributes(Sequence):
         "string8":  np.dtype('S8'),
     }
 
-    def __init__(self, field, segyfd, traceheader_layout, tracecount):
-        super(Attributes, self).__init__(tracecount)
+    def __init__(self, segyfile, field, traceheader_index):
+        super(Attributes, self).__init__(segyfile.tracecount)
         self.field = field
-        self.segyfd = segyfd
-        self.tracecount = tracecount
+        self.traceheader_index = traceheader_index
+        self.segyfd = segyfile.segyfd
+        self.tracecount = segyfile.tracecount
 
+        traceheader_layout = list(segyfile._traceheader_layouts.values())[traceheader_index]
         entry = traceheader_layout.entry_by_byte(field)
         self.dtype = Attributes.ENTRY_TYPE_TO_NUMPY[entry.type]
 
@@ -1109,7 +1111,7 @@ def __getitem__(self, i):
             xs = np.asarray(i, dtype=np.int32)
             xs = xs.astype(dtype=np.int32, order='C', copy=False)
             attrs = np.empty(len(xs), dtype = self.dtype)
-            return self.segyfd.field_foreach(attrs, 0, xs, self.field)
+            return self.segyfd.field_foreach(attrs, self.traceheader_index, xs, self.field)
 
         except TypeError:
             try:
@@ -1124,7 +1126,7 @@ def __getitem__(self, i):
             start, stop, step = i.indices(traces)
             indices = range(start, stop, step)
             attrs = np.empty(len(indices), dtype = self.dtype)
-            return segyfd.field_forall(attrs, 0, start, stop, step, field)
+            return segyfd.field_forall(attrs, self.traceheader_index, start, stop, step, field)
 
 class Text(Sequence):
     """Interact with segy in text mode
@@ -1596,3 +1598,27 @@ def __str__(self):
         type = self.type()
         non_zero = self.use_only_if_non_zero()
         return f"{name} (offset={offset}, type={type}, use_only_if_non_zero={non_zero})"
+
+    def __getitem__(self, key):
+        """ tracefield[key]
+
+        Returns values of traces provided in 'key'.
+        Wrapper around :class:`.Attributes`, so refer to it for more
+        information.
+
+        Parameters
+        ----------
+        key : int or slice
+            Index or slice.
+
+        Returns
+        -------
+        array of numbers//bytearray
+
+        Examples
+        --------
+        Get the cdp_x attribute from trace header extension 1 for first 10 traces:
+
+        >>> f.tracefield.SEG00001.cdp_x[0:10]
+        """
+        return Attributes(self.segyfile, self.entry.byte, self.traceheader_index)[key]

python/test/segy.py

@@ -1079,6 +1079,19 @@ def test_traces_traceheaders(small):
             assert f.traceheader[2][i] == f.traceheader[f.tracecount - 3][i]
 
 
+def test_traceheader_attributes(small):
+    with segyio.open(small, "r") as f:
+        attrs = f.attributes(f.tracefield.SEG00000.xline.offset())
+        assert len(attrs) == f.tracecount
+        assert np.array_equal(list(attrs), [20, 21, 22, 23, 24] * 5)
+
+        assert np.array_equal(f.tracefield.SEG00000.xline[:], attrs)
+        assert np.array_equal(f.tracefield.SEG00000.xline[0:5], attrs[0:5])
+        assert np.array_equal(f.tracefield.SEG00000.xline[0, 5, 11, 12], attrs[0, 5, 11, 12])
+
+        assert f.tracefield.SEG00000.xline[0] == 20
+
+
 def test_depricated_fields(small):
     with segyio.open(small, "r") as f:
         assert f.bin[BinField.EnsembleTraces] == 25
@@ -2487,7 +2500,7 @@ class HeaderTest:
         opened: numbers.Real
         changed: numbers.Real
 
-    trace0_ext_header1_cases = [
+    trace0_proprietary_cases = [
         HeaderTest("int2", 1, -101),
         HeaderTest("int4", 2, -102),
         HeaderTest("int8", 3, -103),
@@ -2512,7 +2525,7 @@ class HeaderTest:
         HeaderTest("header_name", b"TYPES\x00\x00\x00", b"TYPES   "),
     ]
 
-    trace1_ext_header1_cases = [
+    trace1_proprietary_cases = [
         HeaderTest("int2", -1001, 10001),
         HeaderTest("int4", -1002, 2000002),
         HeaderTest("int8", -1003, 3000000003),
@@ -2537,65 +2550,35 @@ class HeaderTest:
     ]
 
     headers = {
-        0: trace0_ext_header1_cases,
-        1: trace1_ext_header1_cases,
+        0: trace0_proprietary_cases,
+        1: trace1_proprietary_cases,
     }
 
     traceheader_index = 1
 
-    # all internal calls are temporary!
-    # they would be exchanged with proper ones once python interface is in place
     with segyio.open(tmpdir / filename, mode='r+', endian=endianness) as f:
-        def mkempty():
-            return bytearray(240)
-
-        types_layout = f._traceheader_layouts["TYPES"]
+        types = f.tracefield[traceheader_index]
         for trace_index, header_cases in headers.items():
             for case in header_cases:
-                offset = types_layout.entry_by_name(case.name).byte
+                offset = types.__getattr__(case.name).offset()
 
-                # assert f.traceheader[trace_index][traceheader_index][offset] == case.opened
-                traceheader = f.segyfd.getth(
-                    trace_index, traceheader_index, mkempty())
-                assert f.segyfd.getfield(
-                    traceheader, traceheader_index, offset) == case.opened
-
-                # f.traceheader[trace_index][traceheader_index][offset] = case.changed
-                f.segyfd.putfield(
-                    traceheader, traceheader_index, offset, case.changed)
-                f.segyfd.putth(trace_index, traceheader_index, traceheader)
-
-                # assert f.traceheader[trace_index][traceheader_index][offset] == case.changed
-                traceheader = f.segyfd.getth(
-                    trace_index, traceheader_index, mkempty())
-                assert f.segyfd.getfield(
-                    traceheader, traceheader_index, offset) == case.changed
+                assert f.traceheader[trace_index][traceheader_index][offset] == case.opened
+                f.traceheader[trace_index][traceheader_index][offset] = case.changed
+                assert f.traceheader[trace_index][traceheader_index][offset] == case.changed
 
         with pytest.raises(ValueError, match=r"Value out of range*"):
-            offset = types_layout.entry_by_name("int2").byte
-
-            # f.traceheader[trace_index][traceheader_index][offset] = 50000
-            f.segyfd.putfield(traceheader, traceheader_index, offset, 50000)
+            offset = types.int2.offset()
+            f.traceheader[trace_index][traceheader_index][offset] = 50000
 
-        from segyio.trace import Attributes
-        for i, case in enumerate(trace0_ext_header1_cases):
+        for i, case in enumerate(trace0_proprietary_cases):
             if (case.name == "scale6_mant") or (case.name == "scale6_exp"):
                 # scale6 is not properly supported yet
                 continue
 
-            if case.name == "header_name":
-                type = "string8"
-            else:
-                type = case.name
-
-            offset = types_layout.entry_by_name(case.name).byte
             expected = np.array(
-                [trace0_ext_header1_cases[i].changed, trace1_ext_header1_cases[i].changed])
-
-            # attrs = f.attributes(offset, traceheader_index)
-            attrs = np.empty(2, dtype=Attributes.ENTRY_TYPE_TO_NUMPY[type])
-            f.segyfd.field_forall(attrs, traceheader_index, 0, 2, 1, offset)
+                [trace0_proprietary_cases[i].changed, trace1_proprietary_cases[i].changed])
 
+            attrs = types.__getattr__(case.name)[:]
             np.testing.assert_array_equal(attrs, expected)