[segyio] 316/376: Add dict-like behaviour for Field

[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 f69efc15a839d8f8187296562150ddef31bfc9eb
Author: Jørgen Kvalsvik <jokva at statoil.com>
Date:   Tue Jun 27 13:20:40 2017 +0200

    Add dict-like behaviour for Field
    
    Supporting the most common dict operation make sense over header
    objects, since they're mappings anyway (order really doesn't matter),
    but they're considered read-only so del and missing are not implemented.
---
 python/segyio/_field.py | 59 ++++++++++++++++++++++++++++++++++++-------------
 python/test/segy.py     | 17 ++++++++++++++
 2 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/python/segyio/_field.py b/python/segyio/_field.py
index d3bca66..0024b13 100644
--- a/python/segyio/_field.py
+++ b/python/segyio/_field.py
@@ -2,23 +2,35 @@ import segyio
 from segyio import BinField
 from segyio import TraceField
 
+class Field(object):
+    _bin_keys = [x for x in BinField.enums()
+                 if  x != BinField.Unassigned1
+                 and x != BinField.Unassigned2]
 
-class Field:
-    def __init__(self, buf, write, field_type, traceno=None):
+    _tr_keys = [x for x in TraceField.enums()
+                if  x != TraceField.UnassignedInt1
+                and x != TraceField.UnassignedInt2]
+
+
+    def __init__(self, buf, write, field_type, traceno=None, keys = _tr_keys):
         self.buf = buf
         self.traceno = traceno
-        self._get_field = segyio._segyio.get_field
-        self._set_field = segyio._segyio.set_field
         self._field_type = field_type
+        self._keys = keys
         self._write = write
 
-    def __getitem__(self, field):
+    def _get_field(self, *args):
+        return segyio._segyio.get_field(self.buf, *args)
 
+    def _set_field(self, *args):
+        return segyio._segyio.set_field(self.buf, *args)
+
+    def __getitem__(self, field):
         # add some structure so we can always iterate over fields
         if isinstance(field, int) or isinstance(field, self._field_type):
             field = [field]
 
-        d = {self._field_type(f): self._get_field(self.buf, f) for f in field}
+        d = {self._field_type(f): self._get_field(f) for f in field}
 
         # unpack the dictionary. if header[field] is requested, a
         # plain, unstructed output is expected, but header[f1,f2,f3]
@@ -28,15 +40,33 @@ class Field:
 
         return d
 
+    def keys(self):
+        return list(self._keys)
+
+    def values(self):
+        return map(self._get_field, self.keys())
+
+    def items(self):
+        return zip(self.keys(), self.values())
+
+    def __contains__(self, key):
+        return key in self._keys
+
+    def __len__(self):
+        return len(self._keys)
+
+    def __iter__(self):
+        return iter(self._keys)
+
     def __setitem__(self, field, val):
-        self._set_field(self.buf, field, val)
+        self._set_field(field, val)
         self._write(self.buf, self.traceno)
 
     def update(self, value):
         buf = self.buf
         if isinstance(value, dict):
             for k, v in value.items():
-                self._set_field(buf, int(k), v)
+                self._set_field(int(k), v)
         else:
             buf = value.buf
 
@@ -56,7 +86,7 @@ class Field:
         def wr(buf, *_):
             segyio._segyio.write_binaryheader(segy.xfd, buf)
 
-        return Field(buf, write=wr, field_type=BinField)
+        return Field(buf, write=wr, field_type=BinField, keys = Field._bin_keys)
 
     @classmethod
     def trace(cls, buf, traceno, segy):
@@ -78,11 +108,10 @@ class Field:
         def wr(buf, traceno):
             segyio._segyio.write_traceheader(segy.xfd, traceno, buf, segy._tr0, segy._bsz)
 
-        return Field(buf, traceno=traceno, write=wr, field_type=TraceField)
+        return Field(buf, traceno=traceno,
+                          write=wr,
+                          field_type=TraceField,
+                          keys=Field._tr_keys)
 
     def __repr__(self):
-        def assigned(x):
-            return x != BinField.Unassigned1 and x != BinField.Unassigned2
-
-        fields = filter(assigned, self._field_type.enums())
-        return self[fields].__repr__()
+        return self[self.keys()].__repr__()
diff --git a/python/test/segy.py b/python/test/segy.py
index c818728..5bc7220 100644
--- a/python/test/segy.py
+++ b/python/test/segy.py
@@ -226,6 +226,23 @@ 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_header_dict_methods(self):
+        with segyio.open(self.filename, "r") as f:
+            self.assertEqual(89, len(list(f.header[0].keys())))
+            self.assertEqual(89, len(list(f.header[1].values())))
+            self.assertEqual(89, len(list(f.header[2].items())))
+            self.assertEqual(89, len(list(f.header[3])))
+            self.assertTrue(0 not in f.header[0])
+            self.assertTrue(1 in f.header[0])
+            self.assertTrue(segyio.su.cdpx in f.header[0])
+            iter(f.header[0])
+
+            self.assertEqual(30, len(f.bin.keys()))
+            self.assertEqual(30, len(list(f.bin.values())))
+            self.assertEqual(30, len(list(f.bin.items())))
+            self.assertEqual(30, len(f.bin))
+            iter(f.bin)
+
     def test_headers_line_offset(self):
         with TestContext("headers_line_offset") as context:
             context.copy_file(self.prestack)

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