Repository: planetarypy/pvl

Issue:
Inserting key/value pairs in PDS label
Is there a way to insert desired key/value lines at a specific point in a PDS label?

ex:
from
```
PDS_VERSION_ID = 
DD_VERSION_ID = 
LABEL_REVISION_NOTE =
^IMAGE
```

to
```
PDS_VERSION_ID = 
DD_VERSION_ID = 
LABEL_REVISION_NOTE =
DATA_SET_ID =
PRODUCT_ID = 
INSTRUMENT_HOST_NAME = 
^IMAGE
```

Discussion:

Patch:
diff --git a/pvl/_collections.py b/pvl/_collections.py
index 2abaf3a..dcb4b91 100644
--- a/pvl/_collections.py
+++ b/pvl/_collections.py
@@ -41,6 +41,10 @@ class KeysView(MappingView):
         keys = [key for key, _ in self._mapping]
         return '%s(%r)' % (type(self).__name__, keys)
 
+    def index(self, key):
+        keys = [k for k, _ in self._mapping]
+        return keys.index(key)
+
 
 class ItemsView(MappingView):
     def __contains__(self, item):
@@ -54,6 +58,10 @@ class ItemsView(MappingView):
     def __getitem__(self, index):
         return self._mapping[index]
 
+    def index(self, item):
+        items = [i for i in self._mapping]
+        return items.index(item)
+
 
 class ValuesView(MappingView):
     def __contains__(self, value):
@@ -73,6 +81,10 @@ class ValuesView(MappingView):
         values = [value for _, value in self._mapping]
         return '%s(%r)' % (type(self).__name__, values)
 
+    def index(self, value):
+        values = [val for _, val in self._mapping]
+        return values.index(value)
+
 
 class OrderedMultiDict(dict, MutableMapping):
     """A ``dict`` like container.
@@ -250,6 +262,67 @@ class OrderedMultiDict(dict, MutableMapping):
     def copy(self):
         return type(self)(self)
 
+    def __insert_wrapper(func):
+        """Make sure the arguments given to the insert methods are correct"""
+        def check_func(self, key, new_item, instance=0):
+            if key not in self.keys():
+                raise KeyError("%s not a key in label" % (key))
+            if not isinstance(new_item, (list, OrderedMultiDict)):
+                raise TypeError("The new item must be a list or PVLModule")
+            if isinstance(new_item, OrderedMultiDict):
+                new_item = list(new_item)
+            return func(self, key, new_item, instance)
+        return check_func
+
+    def _get_index_for_insert(self, key, instance):
+        """Get the index of the key to insert before or after"""
+        if instance == 0:
+            # Index method will return the first occurence of the key
+            index = self.keys().index(key)
+        else:
+            occurrence = -1
+            for index, k in enumerate(self.keys()):
+                if k == key:
+                    occurrence += 1
+                    if occurrence == instance:
+                        # Found the key and the correct occurence of the key
+                        break
+
+            if occurrence != instance:
+                # Gone through the entire list of keys and the instance number
+                # given is too high for the number of occurences of the key
+                raise ValueError(
+                    (
+                        "Cannot insert before/after the %d "
+                        "instance of the key '%s' since there are "
+                        "only %d occurences of the key" % (
+                            instance, key, occurrence)
+                    ))
+        return index
+
+    def _insert_item(self, key, new_item, instance, is_after):
+        """Insert a new item before or after another item"""
+        index = self._get_index_for_insert(key, instance)
+        index = index + 1 if is_after else index
+        self.__items = self.__items[:index] + new_item + self.__items[index:]
+        # Make sure indexing works with new items
+        for new_key, new_value in new_item:
+            if new_key in self:
+                value_list = [val for k, val in self.__items if k == new_key]
+                dict_setitem(self, new_key, value_list)
+            else:
+                dict_setitem(self, new_key, [new_value])
+
+    @__insert_wrapper
+    def insert_after(self, key, new_item, instance=0):
+        """Insert an item after a key"""
+        self._insert_item(key, new_item, instance, True)
+
+    @__insert_wrapper
+    def insert_before(self, key, new_item, instance=0):
+        """Insert an item before a key"""
+        self._insert_item(key, new_item, instance, False)
+
 
 class PVLModule(OrderedMultiDict):
     pass


Test patch:
diff --git a/tests/test_collections.py b/tests/test_collections.py
index 50df73d..18b5a4d 100644
--- a/tests/test_collections.py
+++ b/tests/test_collections.py
@@ -358,14 +358,25 @@ def test_py2_items():
         ('a', 3),
     ])
 
-    assert isinstance(module.items(), list)
-    assert module.items() == [('a', 1), ('b', 2), ('a', 3)]
-
-    assert isinstance(module.keys(), list)
-    assert module.keys() == ['a', 'b', 'a']
-
-    assert isinstance(module.values(), list)
-    assert module.values() == [1, 2, 3]
+    items = module.items()
+    assert isinstance(items, list)
+    assert items == [('a', 1), ('b', 2), ('a', 3)]
+    assert items.index(('a', 1)) == 0
+    assert items.index(('b', 2)) == 1
+    assert items.index(('a', 3)) == 2
+
+    keys = module.keys()
+    assert isinstance(keys, list)
+    assert keys == ['a', 'b', 'a']
+    assert keys.index('a') == 0
+    assert keys.index('b') == 1
+
+    values = module.values()
+    assert isinstance(values, list)
+    assert values == [1, 2, 3]
+    assert values.index(1) == 0
+    assert values.index(2) == 1
+    assert values.index(3) == 2
 
 
 @pytest.mark.skipif(six.PY2, reason='requires python3')
@@ -391,14 +402,30 @@ def test_py3_items():
     ])
 
     assert isinstance(module.items(), pvl._collections.ItemsView)
-    assert module.items()[0] == ('a', 1)
+    items = module.items()
+    assert items[0] == ('a', 1)
+    assert items[1] == ('b', 2)
+    assert items[2] == ('a', 3)
+    assert items.index(('a', 1)) == 0
+    assert items.index(('b', 2)) == 1
+    assert items.index(('a', 3)) == 2
 
     assert isinstance(module.keys(), pvl._collections.KeysView)
-    assert module.keys()[0] == 'a'
+    keys = module.keys()
+    assert keys[0] == 'a'
+    assert keys[1] == 'b'
+    assert keys[2] == 'a'
+    assert keys.index('a') == 0
+    assert keys.index('b') == 1
 
     assert isinstance(module.values(), pvl._collections.ValuesView)
-    assert module.values()[0] == 1
-
+    values = module.values()
+    assert values[0] == 1
+    assert values[1] == 2
+    assert values[2] == 3
+    assert values.index(1) == 0
+    assert values.index(2) == 1
+    assert values.index(3) == 2
 
 
 if six.PY3:
@@ -549,3 +576,153 @@ def test_conversion():
 
     assert dict(module) == expected_dict
     assert list(module) == expected_list
+
+
+@pytest.mark.parametrize(
+    'expected_label, key, instance, expected_list, expected_value', [
+        ([
+            ('a', 4),
+            ('a', 1),
+            ('b', 2),
+            ('a', 3),
+            ('c', 5),
+        ], 'a', 0, [4, 1, 3], 4),
+        ([
+            ('a', 1),
+            ('a', 4),
+            ('b', 2),
+            ('a', 3),
+            ('c', 5),
+        ], 'b', 0, [1, 4, 3], 1),
+        ([
+            ('a', 1),
+            ('b', 2),
+            ('a', 4),
+            ('a', 3),
+            ('c', 5),
+        ], 'a', 1, [1, 4, 3], 1),
+        ([
+            ('a', 1),
+            ('b', 2),
+            ('a', 3),
+            ('a', 4),
+            ('c', 5),
+        ], 'c', 0, [1, 3, 4], 1)
+    ])
+def test_insert_before(expected_label, key, instance, expected_list,
+                        expected_value):
+    module1 = pvl.PVLModule([
+        ('a', 1),
+        ('b', 2),
+        ('a', 3),
+        ('c', 5),
+    ])
+    module2 = module1.copy()
+
+    expected_module = pvl.PVLModule(expected_label)
+
+    module1.insert_before(key, [('a', 4)], instance)
+    assert expected_module == module1
+    assert module1['a'] == expected_value
+    assert module1.getlist('a') == expected_list
+
+    module2.insert_before(key, pvl.PVLModule([('a', 4)]), instance)
+    assert module2 == expected_module
+    assert module1['a'] == expected_value
+    assert module1.getlist('a') == expected_list
+
+
+@pytest.mark.parametrize(
+    'expected_label, key, instance, expected_list, expected_value', [
+        ([
+            ('a', 1),
+            ('a', 4),
+            ('b', 2),
+            ('a', 3),
+            ('c', 5),
+        ], 'a', 0, [1, 4, 3], 1),
+        ([
+            ('a', 1),
+            ('b', 2),
+            ('a', 4),
+            ('a', 3),
+            ('c', 5),
+        ], 'b', 0, [1, 4, 3], 1),
+        ([
+            ('a', 1),
+            ('b', 2),
+            ('a', 3),
+            ('a', 4),
+            ('c', 5),
+        ], 'a', 1, [1, 3, 4], 1),
+        ([
+            ('a', 1),
+            ('b', 2),
+            ('a', 3),
+            ('c', 5),
+            ('a', 4),
+        ], 'c', 0, [1, 3, 4], 1)
+    ])
+def test_insert_after(expected_label, key, instance, expected_list,
+                        expected_value):
+    module1 = pvl.PVLModule([
+        ('a', 1),
+        ('b', 2),
+        ('a', 3),
+        ('c', 5),
+    ])
+    module2 = module1.copy()
+
+    expected_module = pvl.PVLModule(expected_label)
+
+    module1.insert_after(key, [('a', 4)], instance)
+    assert expected_module == module1
+    assert module1['a'] == expected_value
+    assert module1.getlist('a') == expected_list
+
+    module2.insert_after(key, pvl.PVLModule([('a', 4)]), instance)
+    assert module2 == expected_module
+    assert module1['a'] == expected_value
+    assert module1.getlist('a') == expected_list
+
+
+@pytest.mark.parametrize(
+    'key, instance, expected_index', [
+        ('a', 0, 0),
+        ('b', 0, 1),
+        ('a', 1, 2)
+    ])
+def test_get_index_for_insert(key, instance, expected_index):
+    module = pvl.PVLModule([
+        ('a', 1),
+        ('b', 2),
+        ('a', 3),
+    ])
+
+    module._get_index_for_insert(key, instance) == expected_index
+
+
+def test_insert_raises():
+    module = pvl.PVLModule([
+        ('a', 1),
+        ('b', 2),
+        ('a', 3),
+    ])
+
+    with pytest.raises(KeyError):
+        module.insert_before('error_key', [('foo', 'bar')])
+
+    with pytest.raises(KeyError):
+        module.insert_after('error_key', [('foo', 'bar')])
+
+    with pytest.raises(TypeError):
+        module.insert_before('a', ('foo', 'bar'))
+
+    with pytest.raises(TypeError):
+        module.insert_after('a', ('foo', 'bar'))
+
+    with pytest.raises(ValueError):
+        module.insert_before('a', [('foo', 'bar')], 2)
+
+    with pytest.raises(ValueError):
+        module.insert_after('a', [('foo', 'bar')], 2)
diff --git a/tests/test_encoder.py b/tests/test_encoder.py
index 1dee8ff..76fbf1e 100644
--- a/tests/test_encoder.py
+++ b/tests/test_encoder.py
@@ -111,3 +111,33 @@ def test_quoated_strings():
 
     encoder = pvl.encoder.PDSLabelEncoder
     assert module == pvl.loads(pvl.dumps(module, cls=encoder))
+
+
+def test_dump_to_file_insert_before():
+    tmpdir = tempfile.mkdtemp()
+
+    try:
+        for filename in PDS_LABELS:
+            label = pvl.load(filename)
+            if os.path.basename(filename) != 'empty.lbl':
+                label.insert_before('PDS_VERSION_ID', [('new', 'item')])
+            tmpfile = os.path.join(tmpdir, os.path.basename(filename))
+            pvl.dump(label, tmpfile)
+            assert label == pvl.load(tmpfile)
+    finally:
+        shutil.rmtree(tmpdir)
+
+
+def test_dump_to_file_insert_after():
+    tmpdir = tempfile.mkdtemp()
+
+    try:
+        for filename in PDS_LABELS:
+            label = pvl.load(filename)
+            if os.path.basename(filename) != 'empty.lbl':
+                label.insert_after('PDS_VERSION_ID', [('new', 'item')])
+            tmpfile = os.path.join(tmpdir, os.path.basename(filename))
+            pvl.dump(label, tmpfile)
+            assert label == pvl.load(tmpfile)
+    finally:
+        shutil.rmtree(tmpdir)