# coding: utf-8 from __future__ import absolute_import, print_function """ stuff to deal with comments and formatting on dict/list/ordereddict/set these are not really related, formatting could be factored out as a separate base """ import copy from collections import MutableSet, Sized, Set # type: ignore from ruamel.yaml.compat import ordereddict, PY2 __all__ = ["CommentedSeq", "CommentedKeySeq", "CommentedMap", "CommentedOrderedMap", "CommentedSet", 'comment_attrib', 'merge_attrib'] comment_attrib = '_yaml_comment' format_attrib = '_yaml_format' line_col_attrib = '_yaml_line_col' anchor_attrib = '_yaml_anchor' merge_attrib = '_yaml_merge' tag_attrib = '_yaml_tag' class Comment(object): # sys.getsize tested the Comment objects, __slots__ makes them bigger # and adding self.end did not matter __slots__ = 'comment', '_items', '_end', '_start', attrib = comment_attrib def __init__(self): self.comment = None # [post, [pre]] # map key (mapping/omap/dict) or index (sequence/list) to a list of # dict: post_key, pre_key, post_value, pre_value # list: pre item, post item self._items = {} # self._start = [] # should not put these on first item self._end = [] # end of document comments def __str__(self): if self._end: end = ',\n end=' + str(self._end) else: end = '' return "Comment(comment={0},\n items={1}{2})".format( self.comment, self._items, end) @property def items(self): return self._items @property def end(self): return self._end @end.setter def end(self, value): self._end = value @property def start(self): return self._start @start.setter def start(self, value): self._start = value # to distinguish key from None def NoComment(): pass class Format(object): __slots__ = '_flow_style', attrib = format_attrib def __init__(self): self._flow_style = None def set_flow_style(self): self._flow_style = True def set_block_style(self): self._flow_style = False def flow_style(self, default=None): """if default (the flow_style) is None, the flow style tacked on to the object explicitly will be taken. If that is None as well the default flow style rules the format down the line, or the type of the constituent values (simple -> flow, map/list -> block)""" if self._flow_style is None: return default return self._flow_style class LineCol(object): attrib = line_col_attrib def __init__(self): self.line = None self.col = None self.data = None def add_kv_line_col(self, key, data): if self.data is None: self.data = {} self.data[key] = data def key(self, k): return self._kv(k, 0, 1) def value(self, k): return self._kv(k, 2, 3) def _kv(self, k, x0, x1): if self.data is None: return None data = self.data[k] return data[x0], data[x1] def item(self, idx): if self.data is None: return None return self.data[idx][0], self.data[idx][1] def add_idx_line_col(self, key, data): if self.data is None: self.data = {} self.data[key] = data class Anchor(object): __slots__ = 'value', 'always_dump', attrib = anchor_attrib def __init__(self): self.value = None self.always_dump = False class Tag(object): """store tag information for roundtripping""" __slots__ = 'value', attrib = tag_attrib def __init__(self): self.value = None class CommentedBase(object): @property def ca(self): if not hasattr(self, Comment.attrib): setattr(self, Comment.attrib, Comment()) return getattr(self, Comment.attrib) def yaml_end_comment_extend(self, comment, clear=False): if comment is None: return if clear or self.ca.end is None: self.ca.end = [] self.ca.end.extend(comment) def yaml_key_comment_extend(self, key, comment, clear=False): l = self.ca._items.setdefault(key, [None, None, None, None]) if clear or l[1] is None: if comment[1] is not None: assert isinstance(comment[1], list) l[1] = comment[1] else: l[1].extend(comment[0]) l[0] = comment[0] def yaml_value_comment_extend(self, key, comment, clear=False): l = self.ca._items.setdefault(key, [None, None, None, None]) if clear or l[3] is None: if comment[1] is not None: assert isinstance(comment[1], list) l[3] = comment[1] else: l[3].extend(comment[0]) l[2] = comment[0] def yaml_set_start_comment(self, comment, indent=0): """overwrites any preceding comment lines on an object expects comment to be without `#` and possible have multiple lines """ from .error import CommentMark from .tokens import CommentToken pre_comments = self._yaml_get_pre_comment() if comment[-1] == '\n': comment = comment[:-1] # strip final newline if there start_mark = CommentMark(indent) for com in comment.split('\n'): pre_comments.append(CommentToken('# ' + com + '\n', start_mark, None)) def yaml_set_comment_before_after_key(self, key, before=None, indent=0, after=None, after_indent=None): """ expects comment (before/after) to be without `#` and possible have multiple lines """ from ruamel.yaml.error import CommentMark from ruamel.yaml.tokens import CommentToken def comment_token(s, mark): # handle empty lines as having no comment return CommentToken(('# ' if s else '') + s + '\n', mark, None) if after_indent is None: after_indent = indent + 2 if before and before[-1] == '\n': before = before[:-1] # strip final newline if there if after and after[-1] == '\n': after = after[:-1] # strip final newline if there start_mark = CommentMark(indent) c = self.ca.items.setdefault(key, [None, [], None, None]) if before: for com in before.split('\n'): c[1].append(comment_token(com, start_mark)) if after: start_mark = CommentMark(after_indent) if c[3] is None: c[3] = [] for com in after.split('\n'): c[3].append(comment_token(com, start_mark)) @property def fa(self): """format attribute set_flow_style()/set_block_style()""" if not hasattr(self, Format.attrib): setattr(self, Format.attrib, Format()) return getattr(self, Format.attrib) def yaml_add_eol_comment(self, comment, key=NoComment, column=None): """ there is a problem as eol comments should start with ' #' (but at the beginning of the line the space doesn't have to be before the #. The column index is for the # mark """ from .tokens import CommentToken from .error import CommentMark if column is None: column = self._yaml_get_column(key) if comment[0] != '#': comment = '# ' + comment if column is None: if comment[0] == '#': comment = ' ' + comment column = 0 start_mark = CommentMark(column) ct = [CommentToken(comment, start_mark, None), None] self._yaml_add_eol_comment(ct, key=key) @property def lc(self): if not hasattr(self, LineCol.attrib): setattr(self, LineCol.attrib, LineCol()) return getattr(self, LineCol.attrib) def _yaml_set_line_col(self, line, col): self.lc.line = line self.lc.col = col def _yaml_set_kv_line_col(self, key, data): self.lc.add_kv_line_col(key, data) def _yaml_set_idx_line_col(self, key, data): self.lc.add_idx_line_col(key, data) @property def anchor(self): if not hasattr(self, Anchor.attrib): setattr(self, Anchor.attrib, Anchor()) return getattr(self, Anchor.attrib) def yaml_anchor(self): if not hasattr(self, Anchor.attrib): return None return self.anchor def yaml_set_anchor(self, value, always_dump=False): self.anchor.value = value self.anchor.always_dump = always_dump @property def tag(self): if not hasattr(self, Tag.attrib): setattr(self, Tag.attrib, Tag()) return getattr(self, Tag.attrib) def yaml_set_tag(self, value): self.tag.value = value def copy_attributes(self, t, deep=False): for a in [Comment.attrib, Format.attrib, LineCol.attrib, Anchor.attrib, Tag.attrib, merge_attrib]: if hasattr(self, a): if deep: setattr(t, a, copy.deepcopy(getattr(self, a))) else: setattr(t, a, getattr(self, a)) class CommentedSeq(list, CommentedBase): __slots__ = Comment.attrib, def _yaml_add_comment(self, comment, key=NoComment): if key is not NoComment: self.yaml_key_comment_extend(key, comment) else: self.ca.comment = comment def _yaml_add_eol_comment(self, comment, key): self._yaml_add_comment(comment, key=key) def _yaml_get_columnX(self, key): return self.ca.items[key][0].start_mark.column def insert(self, idx, val): """the comments after the insertion have to move forward""" list.insert(self, idx, val) for list_index in sorted(self.ca.items, reverse=True): if list_index < idx: break self.ca.items[list_index+1] = self.ca.items.pop(list_index) def pop(self, idx): res = list.pop(self, idx) self.ca.items.pop(idx, None) # might not be there -> default value for list_index in sorted(self.ca.items): if list_index < idx: continue self.ca.items[list_index-1] = self.ca.items.pop(list_index) return res def _yaml_get_column(self, key): column = None sel_idx = None pre, post = key-1, key+1 if pre in self.ca.items: sel_idx = pre elif post in self.ca.items: sel_idx = post else: # self.ca.items is not ordered for row_idx, k1 in enumerate(self): if row_idx >= key: break if row_idx not in self.ca.items: continue sel_idx = row_idx if sel_idx is not None: column = self._yaml_get_columnX(sel_idx) return column def _yaml_get_pre_comment(self): if self.ca.comment is None: pre_comments = [] self.ca.comment = [None, pre_comments] else: pre_comments = self.ca.comment[1] = [] return pre_comments def __deepcopy__(self, memo): res = CommentedSeq() memo[id(self)] = res for k in self: res.append(copy.deepcopy(k)) self.copy_attributes(res, deep=True) return res class CommentedKeySeq(tuple, CommentedBase): """This primarily exists to be able to roundtrip keys that are sequences""" def _yaml_add_comment(self, comment, key=NoComment): if key is not NoComment: self.yaml_key_comment_extend(key, comment) else: self.ca.comment = comment def _yaml_add_eol_comment(self, comment, key): self._yaml_add_comment(comment, key=key) def _yaml_get_columnX(self, key): return self.ca.items[key][0].start_mark.column def _yaml_get_column(self, key): column = None sel_idx = None pre, post = key-1, key+1 if pre in self.ca.items: sel_idx = pre elif post in self.ca.items: sel_idx = post else: # self.ca.items is not ordered for row_idx, k1 in enumerate(self): if row_idx >= key: break if row_idx not in self.ca.items: continue sel_idx = row_idx if sel_idx is not None: column = self._yaml_get_columnX(sel_idx) return column def _yaml_get_pre_comment(self): if self.ca.comment is None: pre_comments = [] self.ca.comment = [None, pre_comments] else: pre_comments = self.ca.comment[1] = [] return pre_comments class CommentedMapView(Sized): __slots__ = '_mapping', def __init__(self, mapping): self._mapping = mapping def __len__(self): count = len(self._mapping) done = [] # list of processed merge items, kept for masking for merged in getattr(self._mapping, merge_attrib, []): for x in merged[1]: if self._mapping._unmerged_contains(x): continue for y in done: if x in y: break else: count += 1 done.append(merged[1]) return count def __repr__(self): return '{0.__class__.__name__}({0._mapping!r})'.format(self) class CommentedMapKeysView(CommentedMapView, Set): __slots__ = () @classmethod def _from_iterable(self, it): return set(it) def __contains__(self, key): return key in self._mapping def __iter__(self): # yield from self._mapping # not in py27, pypy for x in self._mapping: yield x class CommentedMapItemsView(CommentedMapView, Set): __slots__ = () @classmethod def _from_iterable(self, it): return set(it) def __contains__(self, item): key, value = item try: v = self._mapping[key] except KeyError: return False else: return v == value def __iter__(self): for key in self._mapping._keys(): yield (key, self._mapping[key]) class CommentedMapValuesView(CommentedMapView): __slots__ = () def __contains__(self, value): for key in self._mapping: if value == self._mapping[key]: return True return False def __iter__(self): for key in self._mapping: yield self._mapping[key] class CommentedMap(ordereddict, CommentedBase): __slots__ = Comment.attrib, def _yaml_add_comment(self, comment, key=NoComment, value=NoComment): """values is set to key to indicate a value attachment of comment""" if key is not NoComment: self.yaml_key_comment_extend(key, comment) return if value is not NoComment: self.yaml_value_comment_extend(value, comment) else: self.ca.comment = comment def _yaml_add_eol_comment(self, comment, key): """add on the value line, with value specified by the key""" self._yaml_add_comment(comment, value=key) def _yaml_get_columnX(self, key): return self.ca.items[key][2].start_mark.column def _yaml_get_column(self, key): column = None sel_idx = None pre, post, last = None, None, None for x in self: if pre is not None and x != key: post = x break if x == key: pre = last last = x if pre in self.ca.items: sel_idx = pre elif post in self.ca.items: sel_idx = post else: # self.ca.items is not ordered for row_idx, k1 in enumerate(self): if k1 >= key: break if k1 not in self.ca.items: continue sel_idx = k1 if sel_idx is not None: column = self._yaml_get_columnX(sel_idx) return column def _yaml_get_pre_comment(self): if self.ca.comment is None: pre_comments = [] self.ca.comment = [None, pre_comments] else: pre_comments = self.ca.comment[1] = [] return pre_comments def update(self, vals): try: ordereddict.update(self, vals) except TypeError: # probably a dict that is used for x in vals: self[x] = vals[x] def insert(self, pos, key, value, comment=None): """insert key value into given position attach comment if provided """ ordereddict.insert(self, pos, key, value) if comment is not None: self.yaml_add_eol_comment(comment, key=key) def mlget(self, key, default=None, list_ok=False): """multi-level get that expects dicts within dicts""" if not isinstance(key, list): return self.get(key, default) # assume that the key is a list of recursively accessible dicts def get_one_level(key_list, level, d): if not list_ok: assert isinstance(d, dict) if level >= len(key_list): if level > len(key_list): raise IndexError return d[key_list[level-1]] return get_one_level(key_list, level+1, d[key_list[level-1]]) try: return get_one_level(key, 1, self) except KeyError: return default except (TypeError, IndexError): if not list_ok: raise return default def __getitem__(self, key): try: return ordereddict.__getitem__(self, key) except KeyError: for merged in getattr(self, merge_attrib, []): if key in merged[1]: return merged[1][key] raise def _unmerged_contains(self, key): if ordereddict.__contains__(self, key): return True def __contains__(self, key): if ordereddict.__contains__(self, key): return True # this will only work once the mapping/dict is built to completion for merged in getattr(self, merge_attrib, []): if key in merged[1]: return True return False def get(self, key, default=None): try: return self.__getitem__(key) except: return default def __repr__(self): if not hasattr(self, merge_attrib): return ordereddict.__repr__(self) return 'ordereddict(' + repr(list(self._items())) + ')' def non_merged_items(self): for x in ordereddict.__iter__(self): yield x, ordereddict.__getitem__(self, x) def __delitem__(self, key): found = True for merged in getattr(self, merge_attrib, []): try: del merged[1][key] found = True except KeyError: pass try: ordereddict.__delitem__(self, key) except KeyError: if not found: raise def __iter__(self): for x in ordereddict.__iter__(self): yield x done = [] # list of processed merge items, kept for masking for merged in getattr(self, merge_attrib, []): for x in merged[1]: if ordereddict.__contains__(self, x): continue for y in done: if x in y: break else: yield x done.append(merged[1]) def _keys(self): for x in ordereddict.__iter__(self): yield x done = [] # list of processed merge items, kept for masking for merged in getattr(self, merge_attrib, []): for x in merged[1]: if ordereddict.__contains__(self, x): continue for y in done: if x in y: break else: yield x done.append(merged[1]) if PY2: def keys(self): return list(self._keys()) def iterkeys(self): return self._keys() def viewkeys(self): return CommentedMapKeysView(self) else: def keys(self): return CommentedMapKeysView(self) def _values(self): for x in ordereddict.__iter__(self): yield ordereddict.__getitem__(self, x) done = [] # list of processed merge items, kept for masking for merged in getattr(self, merge_attrib, []): for x in merged[1]: if ordereddict.__contains__(self, x): continue for y in done: if x in y: break else: yield ordereddict.__getitem__(merged[1], x) done.append(merged[1]) if PY2: def values(self): return list(self._values()) def itervalues(self): return self._values() def viewvalues(self): return CommentedMapValuesView(self) else: def values(self): return CommentedMapValuesView(self) def _items(self): for x in ordereddict.__iter__(self): yield x, ordereddict.__getitem__(self, x) done = [] # list of processed merge items, kept for masking for merged in getattr(self, merge_attrib, []): for x in merged[1]: if ordereddict.__contains__(self, x): continue for y in done: if x in y: break else: yield x, ordereddict.__getitem__(merged[1], x) done.append(merged[1]) if PY2: def items(self): return list(self._items()) def iteritems(self): return self._items() def viewitems(self): return CommentedMapItemsView(self) else: def items(self): return CommentedMapItemsView(self) @property def merge(self): if not hasattr(self, merge_attrib): setattr(self, merge_attrib, []) return getattr(self, merge_attrib) def add_yaml_merge(self, value): self.merge.extend(value) def __deepcopy__(self, memo): res = CommentedMap() memo[id(self)] = res for k in self: res[k] = copy.deepcopy(self[k]) self.copy_attributes(res, deep=True) return res class CommentedOrderedMap(CommentedMap): __slots__ = Comment.attrib, class CommentedSet(MutableSet, CommentedMap): __slots__ = Comment.attrib, 'odict', def __init__(self, values=None): self.odict = ordereddict() MutableSet.__init__(self) if values is not None: self |= values def add(self, value): """Add an element.""" self.odict[value] = None def discard(self, value): """Remove an element. Do not raise an exception if absent.""" del self.odict[value] def __contains__(self, x): return x in self.odict def __iter__(self): for x in self.odict: yield x def __len__(self): return len(self.odict) def __repr__(self): return 'set({0!r})'.format(self.odict.keys())