io.mpfile: reinsert comments in get_string

materialsproject · Jul 14, 2015 · 6a2db77 · 6a2db77
1 parent 02f1588
commit 6a2db77
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/mpcontribs/io/mpfile.py b/mpcontribs/io/mpfile.py
@@ -97,6 +97,17 @@ def get_string(self):
                 if lines and key == min_indentor:
                     lines.append('')
                 lines.append(make_pair(key, value, sep=sep))
+        for idx_str, comment in self.comments.iteritems():
+            if idx_str[-1] == '*':
+                lines.insert(int(idx_str[:-1]), '#'+comment)
+            else:
+                idx = int(idx_str)
+                line = lines[idx]
+                table_start = ' '.join([get_indentor(1), 'data_'])
+                if table_start in line:
+                    table_name = line[len(table_start):]
+                    line = ' '.join([get_indentor(1), table_name])
+                lines[idx] = ' #'.join([line, comment])
         return '\n'.join(lines).decode('utf-8')
 
     def __repr__(self):