~abentley/bzrtools/bzrtools.dev

from shutil import rmtree

from bzrlib.branch import Branch

from bzrlib.errors import NoWorkingTree, NotLocalUrl, NotBranchError

from bzrlib.workingtree import WorkingTree

from errors import (NotCheckout, UncommittedCheckout, ParentMissingRevisions,

                    NoParent)

def zap(path, remove_branch=False):

    try:

        wt = WorkingTree.open(path)

    except (NoWorkingTree, NotBranchError):

        raise NotCheckout(path)

    tree_base = wt.bzrdir.transport.base

    branch = wt.branch

    branch_base = branch.bzrdir.transport.base

    if tree_base == branch_base:

        raise NotCheckout(path)

    delta = wt.changes_from(wt.basis_tree(), want_unchanged=False)

    if delta.has_changed():

        raise UncommittedCheckout()

    if remove_branch:

        parent_loc = branch.get_parent()

        if parent_loc is None:

            raise NoParent()

        parent = Branch.open(parent_loc)

        p_ancestry = parent.repository.get_ancestry(parent.last_revision())

        if branch.last_revision() not in p_ancestry:

            raise ParentMissingRevisions(branch.get_parent())

    rmtree(path)

    if remove_branch:

        t = branch.bzrdir.transport

        while t.base != branch_base:

            t = t.clone('..')

        t = t.clone('..')

        t.delete_tree('')

def test_suite():

    import os

    from unittest import makeSuite

    from bzrlib.bzrdir import BzrDir, BzrDirMetaFormat1

    from bzrlib.branch import BranchReferenceFormat

    from bzrlib.tests import TestCaseInTempDir

    class TestZap(TestCaseInTempDir):

        def make_checkout(self):

            wt = BzrDir.create_standalone_workingtree('source')

            os.mkdir('checkout')

            checkout = BzrDirMetaFormat1().initialize('checkout')

            BranchReferenceFormat().initialize(checkout, wt.branch)

            return checkout.create_workingtree()

        def make_checkout2(self):

            wt = self.make_checkout()

            wt2 = wt.branch.bzrdir.sprout('source2').open_workingtree()

            os.mkdir('checkout2')

            checkout = BzrDirMetaFormat1().initialize('checkout2')

            BranchReferenceFormat().initialize(checkout, wt2.branch)

            return checkout.create_workingtree()

        def test_is_checkout(self):

            self.assertRaises(NotCheckout, zap, '.')

            wt = BzrDir.create_standalone_workingtree('.')

            self.assertRaises(NotCheckout, zap, '.')

        def test_zap_works(self):

            self.make_checkout()

            self.assertIs(True, os.path.exists('checkout'))

            zap('checkout')

            self.assertIs(False, os.path.exists('checkout'))

            self.assertIs(True, os.path.exists('source'))

	def test_zap_branch(self):

            self.make_checkout2()

            base = WorkingTree.open('checkout').branch.base

            self.assertIs(True, os.path.exists('checkout'))

            self.assertRaises(NoParent, zap, 'checkout', remove_branch=True)

            self.assertIs(True, os.path.exists('checkout'))

            self.assertIs(True, os.path.exists('source'))

            zap('checkout2', remove_branch=True)

            self.assertIs(False, os.path.exists('checkout2'))

            self.assertIs(False, os.path.exists('source2'))

        def test_checks_modified(self):

            checkout = self.make_checkout()

            os.mkdir('checkout/foo')

            checkout.add('foo')

            self.assertRaises(UncommittedCheckout, zap, 'checkout')

            checkout.commit('commit changes to branch')

            zap('checkout')

    return makeSuite(TestZap)