vrt-pack

#! /usr/bin/env python3
# -*- mode: Python; -*-

from argparse import ArgumentTypeError
from itertools import groupby, count
from tempfile import mkstemp
import enum, os, re, sys, traceback # using enum?
import io, os, string, sys

from vrtargslib import version_args
from vrtargslib import BadData, BadCode
from vrtnamelib import isbinnames

def sizetype(text):
    m = re.fullmatch(r'([1-9][0-9]*)(k|M|)', text)
    if m:
        return int(m.group(1)) * { '' : 10 ** 0,
                                   'k' : 10 ** 3,
                                   'M' : 10 ** 6 }[m.group(2)]
    else:
        raise ArgumentTypeError('bad size')

def parsearguments():
    description = '''

    Pack a directory tree of vrt documents (named *.vrt) to
    predictably sized files (named *.vrf) in a new directory tree of
    fragments sequences for processing, respecting the integrity of
    sentence elements and field-name comments. The intention is that
    eventual unpacking reproduces the original hierarchy with any new
    annotations added to the tokens in the packed fragments.

    '''

    parser = version_args(description = description)

    parser.add_argument('indir', help = 'input directory')

    outgrp = parser.add_mutually_exclusive_group(required = True)
    outgrp.add_argument('--out', '-o', dest = 'outdir', metavar = 'outdir',
                        help = 'output directory')
    outgrp.add_argument('--suffix', '-s', metavar = '.fix',
                        help = 'output in indir.fix/')

    # not sure if this belongs to this tool or another tool
    # parser.add_argument('--element', '-e', metavar = 'name',
    #                     help = 'element to keep whole (default sentence)')

    szegrp = parser.add_mutually_exclusive_group()
    szegrp.add_argument('--lines', '-l', metavar = 'number',
                        type = sizetype,
                        help = 'number of lines to reach in each output file')
    szegrp.add_argument('--tokens', '-t', metavar = 'number',
                        type = sizetype,
                        help = 'number of token lines to reach in each output file')
    szegrp.add_argument('--bytes', '-b', metavar = 'number',
                        type = sizetype,
                        help = 'number of bytes to reach in each output file')

    args = parser.parse_args()
    args.prog = parser.prog
    return args

# A dirsource produces all *.vrt files found under a given directory,
# recursively, in a lexicographic order of their pathnames.

def dirsource(path, memberpath = ''):
    '''Yield a fragment producer for each *.vrt file under path in some
    sort of lexicographic order of their path names, memberpath stored in
    each fragment for eventual unpacking.

    '''

    if os.path.isdir(path):
        for name in sorted(os.listdir(path)):
            yield from dirsource(os.path.join(path, name),
                                 os.path.join(memberpath, name))
    elif os.path.isfile(path) and os.path.splitext(path)[1] == '.vrt':
        inf = open(path, mode = 'br')
        yield makesource(properlines(inf), memberpath)
        inf.close()
    else:
        pass

def properlines(lines):
    for line in lines:
        if line.isspace(): continue
        yield (line if line.endswith(b'\n') else line + b'\n')

def makesource(inf, path):
    '''Return a source of the fragments of the given file together with
    any field-name line found in the file (or None if not found yet -
    it should be in the first fragment, but any further fragment may
    end up in a different output file where the names need to be
    copied so that that file also carries the necessary information.)

    '''

    names = None

    # TODO fix any weird characters in source ... (like, sigh) or
    # maybe refuse to deal with any potentially foul play

    begin = ( b'<... fragment="{}" source="[]">\n'
              .replace(b'[]', path.encode('UTF-8')))
    end = b'</...>\n'

    frak = 1
    def identify(line):
        nonlocal frak, names

        if line == end:
            frak += 1
            return frak - 1

        if isbinnames(line):
            names = line

        return frak

    def bracketed(lines):
        k = count(start = 1)

        start = True
        for line in lines:
            if start:
                yield begin.replace(b'{}', str(next(k)).encode('UTF-8'))
                start = False
            yield line
            if line.startswith(b'</sentence>'):
                yield end
                start = True

        if not start:
            yield end

    return ((names, group)
            for kind, group in
            groupby(bracketed(inf), identify))

def membergen():

    '''Generate up to 27 000 000 archive member names from a000/m000.vrf
    to z999/m999.vrf before crashing. Surely that is more than will
    ever be needed? (That is up to 27 000 directories on the top level
    of the archive but only 1 000 files each. More likely in practice
    there will be only a few top-level directories with the 1 000
    files in each, or just one with a few dozen files.)

    '''

    for a in string.ascii_lowercase:
        for ddd in range(1000):
            for eee in range(1000):
                yield '{}{:03}/m{:03}.vrf'.format(a, ddd, eee)

# TODO this should not be global
membernames = membergen()

# A dirsink creates a directory where the packed fragments go in
# actual files in subdirectories.

def dirsink(args, dirobj, sentinel):
    os.mkdir(dirobj)
    return makesink(args, dirobj, dirobj, dirmember, sentinel)

def dirmember(dirobj, dirname, fieldnames):
    membername = next(membernames)

    subdir = os.path.dirname(os.path.join(dirname, membername))
    os.makedirs(subdir, exist_ok = True)
    fd, temp = mkstemp(dir = subdir,
                       prefix = os.path.basename(membername),
                       suffix = '.tmp')
    os.close(fd)
    out = open(temp, mode = 'bw')

    def end():
        out.close()
        os.rename(temp, os.path.join(dirname, membername))

    # must not print anything outside fragment brackets
    # but fieldnames must be printed in first fragment
    # in every member instead - after fragment bracket;
    # fieldnames is None or out.write(fieldnames)
    return out, end

def makesink(args, target, targetname, member, sentinel):
    size_in_units = ( len if args.bytes else
                      (lambda line: 1) if args.lines else
                      (lambda line: 1 - line.startswith(b'<')) )

    size_limit = ( args.bytes or
                   args.lines or
                   args.tokens or
                   100000 )

    out, end, size = None, None, None
    def consumer(fieldnames, fragment):
        nonlocal out, end, size

        if out is None and fragment is sentinel:
            return

        if fragment is sentinel:
            end()
            return

        if out is None:
            out, end = member(target, targetname, fieldnames)
            size = 0

        if size == 0 and fieldnames is not None:
            # insert field names in first fragment of member
            line = next(fragment)
            out.write(line)
            out.write(fieldnames)
            size += size_in_units(line)
            size += size_in_units(fieldnames)
            
        for line in fragment:
            out.write(line)
            size += size_in_units(line)

        if size >= size_limit:
            end()
            out, end = None, None

    return consumer

def main(args):
    try:
        implement_main(args)
    except KeyboardInterrupt:
        print(args.prog + ': keyboard interrupt',
              file = sys.stderr)
        exit(1)

def implement_main(args):
    indir = os.path.realpath(args.indir)
    outdir = ( os.path.realpath(args.outdir)
               if args.outdir
               else os.path.join(os.path.dirname(indir),
                                 os.path.basename(indir) + args.suffix) )

    if os.path.isdir(indir):
        source = dirsource(indir)
    else:
        print('{}: error: not a directory: {}'
              .format(args.prog, args.indir),
              file = sys.stderr)
        exit(1)

    sentinel = [ b'*** sentinel line ***\n' ]
    try:
        sink = dirsink(args, outdir, sentinel)
    except Exception as exn:
        print('{}: error: could not make output directory'
              .format(args.prog),
              '{}: {}'.format(args.prog, exn),
              sep = '\n', file = sys.stderr)
        exit(1)

    for producer in source:
        for names, fragment in producer:
            sink(names, fragment)

    sink(None, sentinel)

if __name__ == '__main__':
    main(parsearguments())