<commit_before>
import msgpack
import pprint
import ast
from schwa import dr
from schwa.dr.constants import FieldType
from drcli.api import App
from drcli.appargs import ArgumentParser, ISTREAM_AP, OSTREAM_AP, DESERIALISE_AP

META_TYPE = 0

class DumpApp(App):
  """
  Debug: unpack the stream and pretty-print it.
  """
  dump_ap = ArgumentParser()
  dump_ap.add_argument('-m', '--human', dest='human_readable', action='store_true', default=False, help='Reinterpret the messages to be more human-readable by integrating headers into content.')
  dump_ap.add_argument('-n', '--numbered', action='store_true', default=False, help='In --human mode, add a \'#\' field to each annotation, indicating its ordinal index')
  arg_parsers = (dump_ap, ISTREAM_AP, OSTREAM_AP)

  def dump(self, obj):
    pprint.pprint(obj, self.args.out_stream)

  def __call__(self):
    unpacker = msgpack.Unpacker(self.args.in_stream)
    if self.args.human_readable:
      unpacker = self._integrate_names(unpacker)
    for obj in unpacker:
      self.dump(obj)

  def _integrate_names(self, unpacker):
    while True:
      obj = {}
      types = unpacker.unpack()
      if types is None:
        # No new header
        break
      elif isinstance(types, int):
        obj['__version__'] = types
        types = unpacker.unpack()
      store_defs = list(self._process_store_defs(unpacker.unpack(), types))
      nbytes = unpacker.unpack()
      obj['__meta__'] = self._process_annot(unpacker.unpack(), types[META_TYPE][1])
      for store_name, store in store_defs:
        nbytes = unpacker.unpack()
        store['items'] = [self._process_annot(item, store['fields']) for item in unpacker.unpack()]
        if self.args.numbered:
          for i, item in enumerate(store['items']):
            item['#'] = i
        store['fields'] = dict(self._fields_to_dict(store['fields'], store_defs))
        # store.pop('fields')
        obj[store_name] = store
      yield obj

  def _process_store_defs(self, msg, types):
    for name, typ, size in msg:
      try:
        type_name, type_fields = types[typ]
      except IndexError:
        # for robustness to broken data
        type_name, type_fields = '??MissingType={0}'.format(typ), ()
      yield name, {'type': type_name, 'fields': type_fields, 'count': size}

  def _process_annot(self, msg, fields):
    return dict((fields[fnum][FIELD_TYPE_NAME], val) for fnum, val in msg.iteritems())

  TRAIT_NAMES = {
    FieldType.IS_SLICE: 'is slice',
    FieldType.IS_SELF_POINTER: 'is self-pointer',
    FieldType.IS_COLLECTION: 'is collection',
  }

  def _fields_to_dict(self, fields, store_defs, trait_names=TRAIT_NAMES):
    for field in fields:
      name = field.pop(FIELD_TYPE_NAME)
      try:
        field['points to'], store_data = store_defs[field.pop(FIELD_TYPE_POINTER_TO)]
      except KeyError:
        pass
      for trait_num, trait_name in trait_names.items():
        try:
          field[trait_name] = field.pop(trait_num)
        except KeyError:
          pass
      yield name, field


class HackHeaderApp(App):
  """
  Debug: rewrite header components of given documents using Python literal input
  """
  hack_ap = ArgumentParser()
  hack_ap.add_argument('--klasses', default=None, help='Overwrites the entire klasses header with the given list')
  hack_ap.add_argument('-k', '--klass', default=[], action='append', help='Overwrites a klass definition, specified with <name|num>=[<new_name>,<field list>]')
  hack_ap.add_argument('-f', '--field', default=[], action='append', help='Overwrites a field definition, specified with <klass-name|num>.<field-name|num>[+]=<map> (use += for update semantics)')
  arg_parsers = (hack_ap, ISTREAM_AP, OSTREAM_AP)

  def __init__(self, argparser, args):
    super(HackHeaderApp, self).__init__(argparser, args)

    def parse(s, exp_type):
      try:
        res = ast.literal_eval(s)
      except (SyntaxError, ValueError):
        argparser.error('{0} is not a valid Python literal'.format(s))
      if exp_type is not None and type(res) != exp_type:
        argparser.error('{0} does not evaluate to type {1}'.format(s, exp_type))
      return res
    
    self.operations = []

    if args.klasses:
      self.operations.append((self._set_klasses, {'value': parse(args.klasses, list)}))

    for arg in args.klass:
      try:
        key, value = arg.split('=', 1)
      except ValueError:
        argparser.error('Expected <name>=<value>, got {0}'.format(arg))
      try:
        key = int(key)
      except ValueError:
        pass
      value = parse(value, list)
      if len(value) != 2:
        argparser.error('Expected a list of length 2, got {0}'.format(value))
      self.operations.append((self._set_klass, {'klass': key, 'value': value}))

    for arg in args.field:
      try:
        key, value = arg.split('=', 1)
        kname, fname = key.split('.')
      except ValueError:
        argparser.error('Expected <kname>.<fname>=<value>, got {0}'.format(arg))
      if fname.endswith('+'):
        fname = fname[:-1]
        update = True
      else:
        update = False
      try:
        kname = int(kname)
      except ValueError:
        pass
      try:
        fname = int(fname)
      except ValueError:
        pass
      value = parse(value, dict)
      self.operations.append((self._set_field, {'klass': kname, 'field': fname, 'value': value, 'update': update}))
    
    if not self.operations:
      argparser.error('Nothing to do!')

  def _set_klasses(self, klasses, stores, value):
    klasses[:] = value

  def _set_klass(self, klasses, stores, klass, value):
    if klass == len(klasses):
      klasses.append(value)
    for knum, (kname, fields) in enumerate(klasses):
      if klass in (knum, kname):
        klasses[knum] = value
        return
    raise ValueError('Could not find class {0}'.format(klass))

  def _set_field(self, klasses, stores, klass, field, value, update=False):
    for knum, (kname, fields) in enumerate(klasses):
      if klass not in (knum, kname):
        continue

      if field == len(fields):
        fields.append({})

      for fnum, fdef in enumerate(fields):
        fname = fdef.get(FIELD_TYPE_NAME)
        if field in (fnum, fname):
          if update:
            fields[fnum].update(value)
          else:
            fields[fnum] = value
          return

    raise ValueError('Could not find field {1} in class {0}'.format(klass, field))
  
  def __call__(self):
    writer = self.raw_stream_writer
    for doc in self.raw_stream_reader:
      for fn, kwargs in self.operations:
        fn(doc.klasses, doc.stores, **kwargs)
      writer.write(doc)


DumpApp.register_name('dump')
HackHeaderApp.register_name('hackheader')
<commit_msg>Fixed update to new FieldType constants<commit_after>
import msgpack
import pprint
import ast
from schwa import dr
from schwa.dr.constants import FieldType
from drcli.api import App
from drcli.appargs import ArgumentParser, ISTREAM_AP, OSTREAM_AP, DESERIALISE_AP

META_TYPE = 0

class DumpApp(App):
  """
  Debug: unpack the stream and pretty-print it.
  """
  dump_ap = ArgumentParser()
  dump_ap.add_argument('-m', '--human', dest='human_readable', action='store_true', default=False, help='Reinterpret the messages to be more human-readable by integrating headers into content.')
  dump_ap.add_argument('-n', '--numbered', action='store_true', default=False, help='In --human mode, add a \'#\' field to each annotation, indicating its ordinal index')
  arg_parsers = (dump_ap, ISTREAM_AP, OSTREAM_AP)

  def dump(self, obj):
    pprint.pprint(obj, self.args.out_stream)

  def __call__(self):
    unpacker = msgpack.Unpacker(self.args.in_stream)
    if self.args.human_readable:
      unpacker = self._integrate_names(unpacker)
    for obj in unpacker:
      self.dump(obj)

  def _integrate_names(self, unpacker):
    while True:
      obj = {}
      types = unpacker.unpack()
      if types is None:
        # No new header
        break
      elif isinstance(types, int):
        obj['__version__'] = types
        types = unpacker.unpack()
      store_defs = list(self._process_store_defs(unpacker.unpack(), types))
      nbytes = unpacker.unpack()
      obj['__meta__'] = self._process_annot(unpacker.unpack(), types[META_TYPE][1])
      for store_name, store in store_defs:
        nbytes = unpacker.unpack()
        store['items'] = [self._process_annot(item, store['fields']) for item in unpacker.unpack()]
        if self.args.numbered:
          for i, item in enumerate(store['items']):
            item['#'] = i
        store['fields'] = dict(self._fields_to_dict(store['fields'], store_defs))
        # store.pop('fields')
        obj[store_name] = store
      yield obj

  def _process_store_defs(self, msg, types):
    for name, typ, size in msg:
      try:
        type_name, type_fields = types[typ]
      except IndexError:
        # for robustness to broken data
        type_name, type_fields = '??MissingType={0}'.format(typ), ()
      yield name, {'type': type_name, 'fields': type_fields, 'count': size}

  def _process_annot(self, msg, fields):
    return dict((fields[fnum][FieldType.NAME], val) for fnum, val in msg.iteritems())

  TRAIT_NAMES = {
    FieldType.IS_SLICE: 'is slice',
    FieldType.IS_SELF_POINTER: 'is self-pointer',
    FieldType.IS_COLLECTION: 'is collection',
  }

  def _fields_to_dict(self, fields, store_defs, trait_names=TRAIT_NAMES):
    for field in fields:
      name = field.pop(FieldType.NAME)
      try:
        field['points to'], store_data = store_defs[field.pop(FieldType.POINTER_TO)]
      except KeyError:
        pass
      for trait_num, trait_name in trait_names.items():
        try:
          field[trait_name] = field.pop(trait_num)
        except KeyError:
          pass
      yield name, field


class HackHeaderApp(App):
  """
  Debug: rewrite header components of given documents using Python literal input
  """
  hack_ap = ArgumentParser()
  hack_ap.add_argument('--klasses', default=None, help='Overwrites the entire klasses header with the given list')
  hack_ap.add_argument('-k', '--klass', default=[], action='append', help='Overwrites a klass definition, specified with <name|num>=[<new_name>,<field list>]')
  hack_ap.add_argument('-f', '--field', default=[], action='append', help='Overwrites a field definition, specified with <klass-name|num>.<field-name|num>[+]=<map> (use += for update semantics)')
  arg_parsers = (hack_ap, ISTREAM_AP, OSTREAM_AP)

  def __init__(self, argparser, args):
    super(HackHeaderApp, self).__init__(argparser, args)

    def parse(s, exp_type):
      try:
        res = ast.literal_eval(s)
      except (SyntaxError, ValueError):
        argparser.error('{0} is not a valid Python literal'.format(s))
      if exp_type is not None and type(res) != exp_type:
        argparser.error('{0} does not evaluate to type {1}'.format(s, exp_type))
      return res
    
    self.operations = []

    if args.klasses:
      self.operations.append((self._set_klasses, {'value': parse(args.klasses, list)}))

    for arg in args.klass:
      try:
        key, value = arg.split('=', 1)
      except ValueError:
        argparser.error('Expected <name>=<value>, got {0}'.format(arg))
      try:
        key = int(key)
      except ValueError:
        pass
      value = parse(value, list)
      if len(value) != 2:
        argparser.error('Expected a list of length 2, got {0}'.format(value))
      self.operations.append((self._set_klass, {'klass': key, 'value': value}))

    for arg in args.field:
      try:
        key, value = arg.split('=', 1)
        kname, fname = key.split('.')
      except ValueError:
        argparser.error('Expected <kname>.<fname>=<value>, got {0}'.format(arg))
      if fname.endswith('+'):
        fname = fname[:-1]
        update = True
      else:
        update = False
      try:
        kname = int(kname)
      except ValueError:
        pass
      try:
        fname = int(fname)
      except ValueError:
        pass
      value = parse(value, dict)
      self.operations.append((self._set_field, {'klass': kname, 'field': fname, 'value': value, 'update': update}))
    
    if not self.operations:
      argparser.error('Nothing to do!')

  def _set_klasses(self, klasses, stores, value):
    klasses[:] = value

  def _set_klass(self, klasses, stores, klass, value):
    if klass == len(klasses):
      klasses.append(value)
    for knum, (kname, fields) in enumerate(klasses):
      if klass in (knum, kname):
        klasses[knum] = value
        return
    raise ValueError('Could not find class {0}'.format(klass))

  def _set_field(self, klasses, stores, klass, field, value, update=False):
    for knum, (kname, fields) in enumerate(klasses):
      if klass not in (knum, kname):
        continue

      if field == len(fields):
        fields.append({})

      for fnum, fdef in enumerate(fields):
        fname = fdef.get(FieldType.NAME)
        if field in (fnum, fname):
          if update:
            fields[fnum].update(value)
          else:
            fields[fnum] = value
          return

    raise ValueError('Could not find field {1} in class {0}'.format(klass, field))
  
  def __call__(self):
    writer = self.raw_stream_writer
    for doc in self.raw_stream_reader:
      for fn, kwargs in self.operations:
        fn(doc.klasses, doc.stores, **kwargs)
      writer.write(doc)


DumpApp.register_name('dump')
HackHeaderApp.register_name('hackheader')

{'fraction_non_alphanumeric': 0.12860970526942542, 'fraction_numerical': 0.00253051503423638, 'mean_word_length': 3.042418772563177, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4794075', 'n_tokens_mistral': 4363, 'n_tokens_neox': 4316, 'n_words': 1117}