Mercurial extension for my local Mozilla repo

myext.py

# - * - coding: UTF-8 - * -

import re
import os
import os.path
import marshal
from mercurial import cmdutil, patch, revset
from mercurial.error import ParseError
from mercurial.node import bin, hex

cmdtable = {}
command = cmdutil.command(cmdtable)

re_bug = re.compile(r"^Bug (\d+) (?:part (\S+))?")

def _is_in_parts(part, parts):
    if not parts:
        # No parts are specified means all
        return True
    if not part:
        # No part is provided
        return False

    part_float = float(part)
    for s in parts.split(","):
        s = s.split("-", 2)
        if len(s) == 1:
            if s[0] == part:
                return True
        elif part_float >= float(s[0]) and part_float <= float(s[1]):
            return True
    return False

def _traverse_draft(repo):
    def traverse(node):
        yield node
        for child in node.children():
            for subnode in traverse(child):
                yield subnode
    for node in traverse(repo["mybase"]):
        yield node

@command("myexport", [], ("hg myexport [BUG] [PARTS]"))
def myexport(ui, repo, bug=None, parts=None):
    # Create outgoing directory
    dir = ui.config("myexport", "dir", "outgoing")
    dir = os.path.join(repo.path, dir)
    if not os.path.isdir(dir):
        os.mkdir(dir)

    def export_node(node):
        match = re_bug.match(node.description())
        if not match:
            return False # Not a patch commit
        if bug and match.group(1) != bug:
            return False # Not the specified bug
        part = match.group(2)
        if not _is_in_parts(part, parts):
            return False # Not in the specified parts

        filename = "bug{}{}.patch".format(
                match.group(1), "-" + part if part else "")
        with open(os.path.join(dir, filename), "wb") as fp:
            cmdutil.export(repo, [node], fp=fp, opts=patch.diffopts(ui))
        ui.status("exported {} as {}\n".format(node, filename))
        return True

    # Export the patch(s)
    exported = False
    if bug:
        for node in _traverse_draft(repo):
            if export_node(node):
                exported = True
    else:
        exported = export_node(repo["."])

    if not exported:
        ui.warn("no patch exported\n")
        return 1

def revset_mybug(repo, subset, x):
    """``mybug(N)```
    Changesets starts with 'Bug N ' in draft.
    """
    err = 'mybug() requires an integer argument.'
    bugstring = revset.getstring(x, err)
    try:
        bug = int(bugstring)
    except Exception:
        raise ParseError(err)

    draft = revset.draft(repo, subset, None)
    prefix = 'Bug {} '.format(bug)
    return [r for r in draft if repo[r].description().startswith(prefix)]

def extsetup(ui):
    revset.symbols['mybug'] = revset_mybug