Added better comment save/restore methods

atrcopy/segments.py

@@ -717,10 +717,39 @@ class DefaultSegment(object):
         comments = []
         for where_index in has_comments:
             raw = self.get_raw_index(indexes[where_index])
-            comment = self.get_comment(indexes[where_index])
+            comment = self.rawdata.extra.comments[rawindex]
             comments.append(comment)
         return has_comments, comments
 
+    def get_comment_restore_data(self, ranges):
+        """Get a chunk of data (designed to be opaque) containing comments,
+        styles & locations that can be used to recreate the comments on an undo
+        """
+        restore_data = []
+        for start, end in ranges:
+            styles = self.style[start:end].copy()
+            items = []
+            for i in range(start, end):
+                rawindex = self.get_raw_index(i)
+                comment = self.rawdata.extra.comments.get(rawindex, "")
+                items.append((rawindex, comment))
+            restore_data.append((start, end, styles, items))
+        return restore_data
+
+    def restore_comments(self, restore_data):
+        """Restore comment styles and data
+        """
+        for start, end, styles, items in restore_data:
+            self.style[start:end] = styles
+            for rawindex, comment in items:
+                if comment:
+                    self.rawdata.extra.comments[rawindex] = comment
+                else:
+                    try:
+                        del self.rawdata.extra.comments[rawindex]
+                    except KeyError:
+                        pass
+
     def get_comments_in_range(self, start, end):
         """Get a list of comments at specified indexes"""
         comments = {}

test/test_segment.py

@@ -261,6 +261,75 @@ class TestComments(object):
         print r
         assert r == [((0, 38), 0), ((38, 48), 1), ((48, 73), 1), ((73, 78), 1), ((78, 118), 2), ((118, 158), 3), ((158, 198), 4), ((198, 200), 5)]
 
+    def test_restore_comments(self):
+        s = self.segment
+        s.set_style_ranges([[0,1000]], data=True)
+        for i in range(0, len(s), 5):
+            s.set_comment([[i,i+1]], "comment at %d" % i)
+
+        s1 = self.segment
+        print len(s1)
+        indexes = [7,12]
+        r = s1.get_comment_restore_data([indexes])
+        print r
+        # force clear comments
+        s1.rawdata.extra.comments = {}
+        s1.style[indexes[0]:indexes[1]] = 0
+        r0 = s1.get_comment_restore_data([indexes])
+        print r0
+        for start, end, style, items in r0:
+            print style
+            assert np.all(style == 0)
+            for rawindex, comment in items:
+                assert not comment
+        s1.restore_comments(r)
+        r1 = s1.get_comment_restore_data([indexes])
+        print r1
+        for item1, item2 in zip(r, r1):
+            print item1
+            print item2
+            for a1, a2 in zip(item1, item2):
+                print a1, a2
+                if hasattr(a1, "shape"):
+                    assert np.all(a1 - a2 == 0)
+                else:
+                    assert a1 == a2
+
+        s2 = self.sub_segment
+        print len(s2)
+        indexes = [5,10]
+        r = s2.get_comment_restore_data([indexes])
+        print r
+        # force clear comments
+        s2.rawdata.extra.comments = {}
+        s2.style[indexes[0]:indexes[1]] = 0
+        r0 = s2.get_comment_restore_data([indexes])
+        print r0
+        for start, end, style, items in r0:
+            print style
+            assert np.all(style == 0)
+            for rawindex, comment in items:
+                assert not comment
+        s2.restore_comments(r)
+        r2 = s2.get_comment_restore_data([indexes])
+        print r2
+        for item1, item2 in zip(r, r2):
+            print item1
+            print item2
+            for a1, a2 in zip(item1, item2):
+                print a1, a2
+                if hasattr(a1, "shape"):
+                    assert np.all(a1 - a2 == 0)
+                else:
+                    assert a1 == a2
+
+        for item1, item2 in zip(r1, r2):
+            print item1
+            print item2
+            # indexes won't be the same, but rawindexes and comments will
+            assert np.all(item1[2] - item2[2] == 0)
+            assert item1[3] == item2[3]
+
 
 if __name__ == "__main__":
     t = TestIndexed()
@@ -275,3 +344,4 @@ if __name__ == "__main__":
     t = TestComments()
     t.setup()
     t.test_split_data_at_comment()
+    t.test_restore_comments()