<commit_before>"""
Apps to get basic statistics and meta-data from documents.
"""
from __future__ import print_function
import sys
from operator import attrgetter
from collections import defaultdict
from drcli.api import App
from drcli.util import read_raw_docs
from drcli.appargs import ArgumentParser, DESERIALISE_AP, DrInputType


def get_store_names(doc):
    return (tup[0] for tup in doc.stores)


class CountFormatter(object):
  AGG_SUM = 'sum'
  AGG_AVG = 'average'
  ALL = 'all'
  FILE = 'file'
  COUNT_BYTES = 'bytes'
  COUNT_ELEMENTS = 'elements'

  def __init__(self, args, out):
    self.args = args
    self.out = out

  def set_fields(self, names):
    pass

  def add_row(self, counts, ind, agg=None, filename=None, unit=COUNT_ELEMENTS):
    pass

  def start(self):
    pass

  def finish(self):
    pass


class CountTableFormatter(CountFormatter):
  def set_fields(self, names):
    if self.args.show_header:
      print(self.args.field_sep.join(names), file=self.out)

  def add_row(self, counts, ind, agg=None, filename=None, unit=CountFormatter.COUNT_ELEMENTS):
    base = self._fmt_counts(counts)
    suffix = None
    if ind == self.FILE:
      suffix = filename
    elif ind == self.ALL:
      if agg == self.AGG_SUM:
        suffix = 'TOTAL'
      elif agg == self.AGG_AVG:
        suffix = 'AVERAGE'
      else:
        raise ValueError('Unknown aggregate for total: %r' % agg)
    print(base + (self.args.field_sep + suffix if suffix else ''), file=self.out)

  def _fmt_counts(self, counts):
    return self.args.field_sep.join(str(c) for c in counts)


class CountJsonFormatter(CountFormatter):
  def __init__(self, *args, **kwargs):
    super(CountJsonFormatter, self).__init__(*args, **kwargs)
    import json
    self._dumps = json.dumps

  def set_fields(self, names):
    self._fields = [self._dumps(name) for name in names]

  def add_row(self, counts, ind, agg=None, filename=None, unit=CountFormatter.COUNT_ELEMENTS):
    if self._row_added:
      self.out.write(',\n')
    self.out.write('{')
    self.out.write(', '.join('%s: %s' % tup for tup in zip(self._fields, counts)))
    self.out.write(', "_meta": {"after": %s' % self._dumps(ind))
    if ind == self.FILE:
        self.out.write(', "file": %s' % self._dumps(filename))
    if agg:
        self.out.write(', "agg": %s' % self._dumps(agg))
    if unit != CountFormatter.COUNT_ELEMENTS:
        self.out.write(', "unit": %s' % self._dumps(unit))
    self.out.write('}}')
    self.out.flush()
    self._row_added = True

  def start(self):
    self._row_added = False
    print('[', file=self.out)

  def finish(self):
    print('\n]', file=self.out)


class CountApp(App):
  """
  Count the number of documents or annotations in named stores.
  """
  # TODO: options to choose stores to count
  # TODO: avoid desiralising?
  count_arg_parser = ArgumentParser()
  count_arg_parser.add_argument('-s', '--store', metavar='ATTR', dest='count_stores', action='append', default=[], help='Count the specified store')
  count_arg_parser.add_argument('-d', '--docs', dest='count_docs', action='store_true', help='Count the number of documents (default without fields specified)')
  count_arg_parser.add_argument('-a', '--all', dest='count_all', action='store_true', help='Count docs and elements in all stores found on the first document')
  count_arg_parser.add_argument('--every', dest='show_interval', type=int, metavar='N', help='Show counts every N docs')
  count_arg_parser.add_argument('-e', '--every1', dest='show_interval', action='store_const', const=1, help='Show counts every doc')
  count_arg_parser.add_argument('--bytes', dest='count_bytes', action='store_true', default=False, help='Count the number of bytes for each store, rather than the number of elements')
  count_arg_parser.add_argument('--no-subtotal', dest='show_subtotal', default=True, action='store_false', help='Hides total count per input file')
  count_arg_parser.add_argument('--no-total', dest='show_total', default=True, action='store_false', help='Hides total count across all documents')
  count_arg_parser.add_argument('--average', dest='show_average', default=False, action='store_true', help='Show an average size per document')
  count_arg_parser.add_argument('--no-header', dest='show_header', default=True, action='store_false', help='Hides the field names displayed by --fmt-table with more than one field output')
  count_arg_parser.add_argument('-c', '--cumulative', default=False, action='store_true', help='Show cumulative counts')
  count_arg_parser.add_argument('--sep', dest='field_sep', default='\t', help='Output field separator (with --fmt-table)')
  count_arg_parser.add_argument('--fmt-table', dest='formatter_cls', action='store_const', const=CountTableFormatter, default=CountTableFormatter, help='Format output as a table (default)')
  count_arg_parser.add_argument('-j', '--fmt-json', dest='formatter_cls', action='store_const', const=CountJsonFormatter, help='Format output as JSON')
  count_arg_parser.add_argument('files', nargs='*', type=DrInputType, help='Specify files by name rather than standard input')
  arg_parsers = (count_arg_parser, ISTREAM_AP,)

  def __init__(self, argparser, args):
    if args.count_all and (args.count_docs or args.count_stores):
      argparser.error('--all flag may not be used in conjunction with --docs or store names')
<commit_msg>Add a dr count option to show timestamp for I/O logging<commit_after>"""
Apps to get basic statistics and meta-data from documents.
"""
from __future__ import print_function
import sys
from operator import attrgetter
from collections import defaultdict
import datetime
from drcli.api import App
from drcli.util import read_raw_docs
from drcli.appargs import ArgumentParser, DESERIALISE_AP, DrInputType


def get_store_names(doc):
    return (tup[0] for tup in doc.stores)

def iso_now():
    return datetime.datetime.now().isoformat()


class CountFormatter(object):
  AGG_SUM = 'sum'
  AGG_AVG = 'average'
  ALL = 'all'
  FILE = 'file'
  COUNT_BYTES = 'bytes'
  COUNT_ELEMENTS = 'elements'

  def __init__(self, args, out):
    self.args = args
    self.out = out

  def set_fields(self, names):
    pass

  def add_row(self, counts, ind, agg=None, filename=None, unit=COUNT_ELEMENTS):
    pass

  def start(self):
    pass

  def finish(self):
    pass


class CountTableFormatter(CountFormatter):
  def set_fields(self, names):
    if self.args.show_header:
      header = self.args.field_sep.join(names)
      if self.args.timestamp:
        header = iso_now() + self.args.field_sep
      print(header, file=self.out)

  def add_row(self, counts, ind, agg=None, filename=None, unit=CountFormatter.COUNT_ELEMENTS):
    base = self._fmt_counts(counts)
    suffix = None
    if ind == self.FILE:
      suffix = filename
    elif ind == self.ALL:
      if agg == self.AGG_SUM:
        suffix = 'TOTAL'
      elif agg == self.AGG_AVG:
        suffix = 'AVERAGE'
      else:
        raise ValueError('Unknown aggregate for total: %r' % agg)
    print(base + (self.args.field_sep + suffix if suffix else ''), file=self.out)

  def _fmt_counts(self, counts):
    res = self.args.field_sep.join(str(c) for c in counts)
    if self.args.timestamp:
      res = iso_now() + self.args.field_sep + res
    return res


class CountJsonFormatter(CountFormatter):
  def __init__(self, *args, **kwargs):
    super(CountJsonFormatter, self).__init__(*args, **kwargs)
    import json
    self._dumps = json.dumps

  def set_fields(self, names):
    self._fields = [self._dumps(name) for name in names]

  def add_row(self, counts, ind, agg=None, filename=None, unit=CountFormatter.COUNT_ELEMENTS):
    if self._row_added:
      self.out.write(',\n')
    self.out.write('{')
    self.out.write(', '.join('%s: %s' % tup for tup in zip(self._fields, counts)))
    self.out.write(', "_meta": {"after": %s' % self._dumps(ind))
    if ind == self.FILE:
      self.out.write(', "file": %s' % self._dumps(filename))
    if agg:
      self.out.write(', "agg": %s' % self._dumps(agg))
    if unit != CountFormatter.COUNT_ELEMENTS:
      self.out.write(', "unit": %s' % self._dumps(unit))
    if self.args.timestamp:
      self.out.write(', "time": %s' % self._dumps(iso_now()))
    self.out.write('}}')
    self.out.flush()
    self._row_added = True

  def start(self):
    self._row_added = False
    print('[', file=self.out)

  def finish(self):
    print('\n]', file=self.out)


class CountApp(App):
  """
  Count the number of documents or annotations in named stores.
  """
  # TODO: options to choose stores to count
  # TODO: avoid desiralising?
  count_arg_parser = ArgumentParser()
  count_arg_parser.add_argument('-s', '--store', metavar='ATTR', dest='count_stores', action='append', default=[], help='Count the specified store')
  count_arg_parser.add_argument('-d', '--docs', dest='count_docs', action='store_true', help='Count the number of documents (default without fields specified)')
  count_arg_parser.add_argument('-a', '--all', dest='count_all', action='store_true', help='Count docs and elements in all stores found on the first document')
  count_arg_parser.add_argument('-v', '--every', dest='show_interval', type=int, metavar='N', help='Show counts every N docs')
  count_arg_parser.add_argument('-e', '--every1', dest='show_interval', action='store_const', const=1, help='Show counts every doc')
  count_arg_parser.add_argument('--bytes', dest='count_bytes', action='store_true', default=False, help='Count the number of bytes for each store, rather than the number of elements')
  count_arg_parser.add_argument('--no-subtotal', dest='show_subtotal', default=True, action='store_false', help='Hides total count per input file')
  count_arg_parser.add_argument('--no-total', dest='show_total', default=True, action='store_false', help='Hides total count across all documents')
  count_arg_parser.add_argument('--average', dest='show_average', default=False, action='store_true', help='Show an average size per document')
  count_arg_parser.add_argument('--no-header', dest='show_header', default=True, action='store_false', help='Hides the field names displayed by --fmt-table with more than one field output')
  count_arg_parser.add_argument('-c', '--cumulative', default=False, action='store_true', help='Show cumulative counts')
  count_arg_parser.add_argument('-t', '--timestamp', action='store_true', default=False, help='Output the time with each count')
  count_arg_parser.add_argument('--sep', dest='field_sep', default='\t', help='Output field separator (with --fmt-table)')
  count_arg_parser.add_argument('--fmt-table', dest='formatter_cls', action='store_const', const=CountTableFormatter, default=CountTableFormatter, help='Format output as a table (default)')
  count_arg_parser.add_argument('-j', '--fmt-json', dest='formatter_cls', action='store_const', const=CountJsonFormatter, help='Format output as JSON')
  count_arg_parser.add_argument('files', nargs='*', type=DrInputType, help='Specify files by name rather than standard input')
  arg_parsers = (count_arg_parser, ISTREAM_AP,)

  def __init__(self, argparser, args):
    if args.count_all and (args.count_docs or args.count_stores):
      argparser.error('--all flag may not be used in conjunction with --docs or store names')

{'fraction_non_alphanumeric': 0.1411837785074097, 'fraction_numerical': 0.0005324341112787292, 'mean_word_length': 4.595829195630586, 'pattern_counts': {'":': 11, '<': 3, '<?xml version=': 0, '>': 3, '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': '210070', 'n_tokens_mistral': 3677, 'n_tokens_neox': 3577, 'n_words': 942}