Ticket #12342: hg_json

#!/usr/bin/env python
"""
Convert between a Mercurial 1.0 repository and a JSON file.

AUTHORS:
  
    - Keshav Kini (Dec 2011)
    - William Stein (Jan 2012)
"""

import os, sys

if not os.environ.has_key('SAGE_ROOT'):
    print "%s: run this using the Sage environment, e.g., type 'sage -sh'"%(sys.argv[0])
    sys.exit(1)

import struct, json, cStringIO, shutil
from mercurial.changegroup import readexactly, readbundle, writebundle

#
# bundle --> JSON
#

def unpack_groups(fh):
    """
    A generator of parsed groups from a bundle file. Decompress the bundle
    and discard the magic header before calling this function (i.e. call this
    function with a "changegroup" as input). Expects bundle format HG10
    (the bundle format introduced in Mercurial 1.0).
    """
    yield [chunk for chunk in unpack_chunks(fh)] # the changeloge group
    yield [chunk for chunk in unpack_chunks(fh)] # the manifest group
    while True:
        length, = struct.unpack('>l', readexactly(fh, 4))
        if length <= 4:
            # found a "null meta chunk", which ends the changegroup
            break
        filename = readexactly(fh, length - 4).encode('string_escape')
        yield (filename, [chunk for chunk in unpack_chunks(fh)]) # a file group

def unpack_chunks(fh):
    """
    A generator of parsed chunks of a "group" in a bundle file. Place the
    input head at the beginning of a group, and this function will yield
    parsed chunks until the end of the group (a "null chunk").
    """
    while True:
        length, = struct.unpack('>l', readexactly(fh, 4))
        if length <= 4:
            # found a "null chunk", which ends the group
            break
        if length < 84:
            raise Exception("negative data length")
        node, p1, p2, cs = struct.unpack( '20s20s20s20s',
                readexactly(fh, 80) )
        yield { 'node': node.encode('hex')
              , 'p1': p1.encode('hex')
              , 'p2': p2.encode('hex')
              , 'cs': cs.encode('hex')
              , 'data': [patch for patch in unpack_patches(fh, length - 84)] }

def unpack_patches(fh, remaining):
    """
    A generator of patches from the data field in a chunk. As there is
    no delimiter for this data field, we require a length argument.
    """
    while remaining >= 12:
        start, end, blocklen = struct.unpack('>lll', readexactly(fh, 12))
        remaining -= 12
        if blocklen > remaining:
            raise Exception("unexpected end of patch stream")
        block = readexactly(fh, blocklen)
        remaining -= blocklen
        yield { 'start': start
              , 'end': end
              , 'blocklen': blocklen
              , 'block': block.encode('string_escape')
              }

    if remaining > 0:
        print remaining
        raise Exception("unexpected end of patch stream")

def to_json(ifilename, ofilename):
    """
    Given an HG10xx file (Mercurial 1.0 bundle) ``ifilename``, convert
    it to JSON and dump it to ``ofilename``
    """
    if ifilename:
        ifile = open(ifilename, 'rb')
    else:
        ifile = sys.stdin
    fh = readbundle(ifile, ifilename)

    oobj = [group for group in unpack_groups(fh)]
    # this is unimplemented in Mercurial 1.8.4, unfortunately (?)
    #fh.close()

    if ofilename:
        ofile = open(ofilename, 'w')
    else:
        ofile = sys.stdout
    json.dump(oobj, ofile, indent=4)
    ofile.close()

#
# JSON --> bundle
#

def pack_groups(cg_obj):
    """
    Given a JSON object ``cg_obj`` representing a changegroup, return the
    changegroup.
    """
    return ''.join([ pack_chunks(cg_obj[0])   # the changelog group
                   , pack_chunks(cg_obj[1]) ] # the manifest group
                 + [ ''.join([ struct.pack('>l', len(pair[0].decode('string_escape')) + 4)
                             , pair[0].decode('string_escape') # a filename
                             , pack_chunks(pair[1]) ]) # a file group
                   for pair in cg_obj[2:] ] # over all files
                 + [ struct.pack('>l', 0) ]) # a null meta chunk to end files

def pack_chunks(group_obj):
    """
    Given a JSON object ``group_obj`` representing a group, create the group
    and write it to ``fh``
    """
    def chunk_gen():
        for chunk in group_obj:
            chunk['data'] = pack_patches(chunk['data'])
            yield ''.join([ struct.pack('>l', len(chunk['data']) + 84)
                          , chunk['node'].decode('string_escape').decode('hex')
                          , chunk['p1'].decode('string_escape').decode('hex')
                          , chunk['p2'].decode('string_escape').decode('hex')
                          , chunk['cs'].decode('string_escape').decode('hex')
                          , chunk['data'] ])
    return ''.join([ packed_chunk for packed_chunk in chunk_gen() ]
                 + [ struct.pack('>l', 0) ]) # a null chunk to end the group

def pack_patches(data_obj):
    """
    Pack a given list of patch objects into a binary patch stream.
    """
    return ''.join([ ''.join([ struct.pack( '>lll'
                                          , patch['start']
                                          , patch['end']
                                          , patch['blocklen'] )
                             , patch['block'].decode('string_escape') ])
                   for patch in data_obj ])

def to_hg(ifilename, ofilename):
    """
    Given a JSON file produced by this script, at ``ifilename``, convert
    it into an HG10UN file (Mercurial 1.0 uncompressed bundle) at
    ``ofilename``.
    """
    if ifilename:
        ifile = open(ifilename, 'r')
    else:
        ifile = sys.stdin
    iobj = json.load(ifile, encoding='ascii')

    writebundle(cStringIO.StringIO(pack_groups(iobj)), ofilename, 'HG10UN')

def system(cmd):
    """
    Verbosely run the given line under the shell.  Raise a
    RuntimeError if the command fails or is killed.
    """
    print cmd
    if os.system(cmd):
        raise RuntimeError, "error running command '%s'"%cmd

def convert_hg_to_json(path, destructive=False):
    """
    Convert path/.hg to path/.hg.json

    Do no error checking that the path contains .hg and not .hg.json;
    that was assumed done by the convert function, which calls this.
    """
    print "Converting .hg repo to .hg.json plaintext file..."
    bundle = os.path.join(path, '.hg_json.bundle')
    try:
        # 1. convert repository to a bundle
        cmd = 'hg bundle -a -R "%s" "%s"'%(path, bundle)
        system(cmd)
        # 2. convert bundle to json dump
        out_json = os.path.join(path, '.hg.json')
        to_json(bundle, out_json)
    finally:
        # 3. Clean up no matter what
        if os.path.exists(bundle):
            os.remove(bundle)
    if not os.path.exists(out_json):
        raise RuntimeError, "something went wrong creating .hg.json"
    print "created '%s'"%out_json
    # 4. Destroy the .hg repository
    hg_repo = os.path.join(path, '.hg')
    if destructive and os.path.exists(hg_repo):
        shutil.rmtree(hg_repo)

def convert_json_to_hg(path, destructive=False):
    """
    Convert path/.hg.json to path/.hg

    Do no error checking that the path contains .hg.json and not .hg;
    that was assumed done by the convert function, which calls this.
    """
    print "Converting .hg.json plaintext file to .hg repo..."
    json_file = os.path.join(path, '.hg.json')
    bundle_file = os.path.join(path, '.hg.json.bundle')
    try:
        # 1. convert json to bundle
        to_hg(json_file, bundle_file)
        if not os.path.exists(bundle_file):
            raise RuntimeError
        # 2. convert bundle to .hg repo
        system('hg init "%s"'%path)
        system('hg pull -R "%s" "%s"'%(path, bundle_file))
        system('hg update -R "%s"'%path)
    finally:
        # 3. Clean up
        if os.path.exists(bundle_file):
            os.remove(bundle_file)
    print "created '%s'"%os.path.join(path, '.hg')
    # 4. Destroy plain text file
    if destructive and os.path.exists(json_file):
        os.remove(json_file)

def convert(path, destructive=False):
    """
    Do some error checking and decide whether we need to convert
    from .hg to .hg.json or conversely, then call the appropriate
    conversion function.
    """
    if not os.path.exists(path):
        raise RuntimeError, "path ('%s') does not exist"%path
    if not os.path.isdir(path):
        raise RuntimeError, "path ('%s') must be a directory"%path
    path_to_hg = os.path.join(path, '.hg')
    path_to_json = os.path.join(path, '.hg.json')
    hg_to_json = os.path.exists(path_to_hg)
    json_to_hg = os.path.exists(path_to_json)
    if hg_to_json and json_to_hg:
        raise RuntimeError, "path ('%s') must contain either an .hg repo or .hg.json, but not both (right now it contains both)"%path
    if not hg_to_json and not json_to_hg:
        raise RuntimeError, "path ('%s') must contain an .hg repo or .hg.json file, but contains neither"%path
    if hg_to_json:
        convert_hg_to_json(path, destructive)
    if json_to_hg:
        convert_json_to_hg(path, destructive)
    

#
# Command line usage
#

if __name__ == '__main__':
    import optparse        # so works with python 2.6

    parser = optparse.OptionParser(
        usage = "Usage: %prog path",
        description="Convert Mercurial repo (.hg directory) to .hg.json file or convert back.",
        epilog=
"""If path contain only an .hg repository, then .hg.json is created.
If path contains only a file .hg.json, then the .hg repository is created.
Any other situation is an error.  The default path is the current directory. 
 """)
    
    parser.add_option('--destructive', dest="destructive",
                      help="(default: 'yes') if 'yes' deletes either .hg or .hg.json after creating .hg.json or .hg", default="yes")

    (options, args) = parser.parse_args()
    if len(args) !=1 :
        parser.print_help()
        sys.exit(1)

    destructive =  options.destructive=='yes'   # yes, ==.
    convert(args[0], destructive)