Source Code for Module rekall.yaml_utils

  1  import collections 
  2  import yaml 
  3   
  4   
  5 -class OrderedYamlDict(yaml.YAMLObject, collections.OrderedDict): 
  6      """A class which produces an ordered dict.""" 
  7      yaml_tag = "tag:yaml.org,2002:map" 
  8   
  9      @classmethod 
 10 -    def to_yaml(cls, dumper, data): 
 11          value = [] 
 12          node = yaml.nodes.MappingNode(cls.yaml_tag, value) 
 13          for key, item in data.iteritems(): 
 14              node_key = dumper.represent_data(key) 
 15              node_value = dumper.represent_data(item) 
 16              value.append((node_key, node_value)) 
 17   
 18          return node 
 19   
 20      @classmethod 
 21 -    def construct_mapping(cls, loader, node, deep=False): 
 22          """Based on yaml.loader.BaseConstructor.construct_mapping.""" 
 23   
 24          if not isinstance(node, yaml.MappingNode): 
 25              raise yaml.loader.ConstructorError( 
 26                  None, None, "expected a mapping node, but found %s" % node.id, 
 27                  node.start_mark) 
 28   
 29          mapping = OrderedYamlDict() 
 30          for key_node, value_node in node.value: 
 31              key = loader.construct_object(key_node, deep=deep) 
 32              try: 
 33                  hash(key) 
 34              except TypeError as exc: 
 35                  raise yaml.loader.ConstructorError( 
 36                      "while constructing a mapping", node.start_mark, 
 37                      "found unacceptable key (%s)" % exc, key_node.start_mark) 
 38   
 39              value = loader.construct_object(value_node, deep=deep) 
 40              mapping[key] = value 
 41   
 42          return mapping 
 43   
 44      @classmethod 
 45 -    def from_yaml(cls, loader, node): 
 46          """Parse the yaml file into an OrderedDict so we can preserve order.""" 
 47          fields = cls.construct_mapping(loader, node, deep=True) 
 48          result = cls() 
 49          for k, v in fields.items(): 
 50              result[k] = v 
 51   
 52          return result 
 53   
 54   
 55 -def decode(data): 
 56      return yaml.safe_load(data) or OrderedYamlDict() 
 57   
 58 -def encode(raw_data): 
 59      return yaml.safe_dump(raw_data, default_flow_style=False) 
 60   
 61   
 62 -class PrettyPrinterDumper(yaml.SafeDumper): 
 63      """A dumper which produces pretty printed YAML. 
 64   
 65      See: 
 66      http://stackoverflow.com/questions/6432605/any-yaml-libraries-in-python-that-support-dumping-of-long-strings-as-block-liter 
 67      """ 
 68   
 69   
 70 -def unicode_representer(_, data): 
 71      has_wide_lines = False 
 72      for line in data.splitlines(): 
 73          if len(line) > 80: 
 74              has_wide_lines = True 
 75              break 
 76   
 77      if has_wide_lines: 
 78          return yaml.ScalarNode( 
 79              u'tag:yaml.org,2002:str', data, style='>') 
 80   
 81      if "\n" in data: 
 82          return yaml.ScalarNode( 
 83              u'tag:yaml.org,2002:str', data, style='|') 
 84   
 85      return yaml.ScalarNode( 
 86          u'tag:yaml.org,2002:str', data, style='') 
 87   
 88 -def represent_orderedyamldict(dumper, data): 
 89      value = [] 
 90   
 91      for item_key, item_value in data.items(): 
 92          node_key = dumper.represent_data(item_key) 
 93          if type(item_value) not in [ 
 94                  str, unicode, list, dict, OrderedYamlDict, bool, long, int]: 
 95              import pdb; pdb.set_trace() 
 96          node_value = dumper.represent_data(item_value) 
 97   
 98          value.append((node_key, node_value)) 
 99   
100      return yaml.nodes.MappingNode(u'tag:yaml.org,2002:map', value) 
101   
102   
103  PrettyPrinterDumper.add_representer( 
104      unicode, unicode_representer) 
105   
106  PrettyPrinterDumper.add_representer( 
107      str, unicode_representer) 
108   
109  PrettyPrinterDumper.add_representer( 
110      OrderedYamlDict, represent_orderedyamldict) 
111   
112   
113 -def safe_dump(data, **kwargs): 
114      kwargs["default_flow_style"] = False 
115      return yaml.dump_all( 
116          [data], None, Dumper=PrettyPrinterDumper, **kwargs) 
117