--- a/hgext/states.py	Wed Sep 28 22:04:52 2011 +0200
+++ b/hgext/states.py	Thu Sep 29 15:23:56 2011 +0200
@@ -343,6 +343,7 @@
 '''
 import os
 from functools import partial
+from operator import or_
 
 from mercurial.i18n import _
 from mercurial import cmdutil
@@ -358,9 +359,11 @@
 from mercurial import wireproto
 from mercurial import pushkey
 from mercurial import error
+from mercurial import repair
 from mercurial.lock import release
 
 
+
 # states property constante
 _NOSHARE=2
 _MUTABLE=1
@@ -714,6 +717,39 @@
                     hint=_('see `hg help --extension states` for details'))
     return orig(queue, repo, *args, **kwargs)
 
+def strip(orig, ui, repo, node, backup="all"):
+    cl = repo.changelog
+    striprev = cl.rev(node)
+    revstostrip = set(cl.descendants(striprev))
+    revstostrip.add(striprev)
+    tostrip = set(map(cl.node, revstostrip))
+    # compute the potentially new created states bondaries which are any
+    # parent of the stripped node that are not stripped (may not be heads)
+    newbondaries = set(par for nod in tostrip for par in cl.parents(nod)
+                       if par not in tostrip)
+    # save the current states of newbondaries in a chache as repo.nodestate
+    # must work along the loop. We will use the next loop to add them.
+    statesheads={}
+    for nd in newbondaries:
+        state = repo.nodestate(nd)
+        if state.trackheads:
+            statesheads.setdefault(state, set([])).add(nd)
+
+    for state, heads in repo._statesheads.iteritems():
+        if not state.trackheads:
+            continue
+        heads = set(heads) - tostrip | statesheads.get(state, set([]))
+        # reduce heads (make them really heads)
+        revs = set(map(cl.rev, heads))
+        minrev = min(revs)
+        for rev in cl.ancestors(*revs):
+            if rev >= minrev:
+                revs.discard(rev)
+        repo._statesheads[state] = map(cl.node, revs)
+    _writestateshead(repo)
+
+    return orig(ui, repo, node, backup)
+
 
 def uisetup(ui):
     """
@@ -723,6 +759,7 @@
     # patch discovery
     extensions.wrapfunction(discovery, 'findcommonoutgoing', filterprivateout)
     extensions.wrapfunction(discovery, 'findcommonincoming', filterprivatein)
+    extensions.wrapfunction(repair, 'strip', strip)
 
     # patch wireprotocol
     wireproto.commands['heads'] = (wireheads, '')
@@ -1032,41 +1069,3 @@
 
     repo.__class__ = statefulrepo
 
-    import mercurial.repair
-    from operator import or_
-    ostrip = mercurial.repair.strip
-
-    def strip(ui, repo, node, backup="all"):
-        cl = repo.changelog
-        striprev = cl.rev(node)
-        revstostrip = set(cl.descendants(striprev))
-        revstostrip.add(striprev)
-        tostrip = set(map(cl.node, revstostrip))
-        # compute the potentially new created states bondaries which are any
-        # parent of the stripped node that are not stripped (may not be heads)
-        newbondaries = set(par for nod in tostrip for par in cl.parents(nod)
-                           if par not in tostrip)
-        # save the current states of newbondaries in a chache as repo.nodestate
-        # must work along the loop. We will use the next loop to add them.
-        statesheads={}
-        for nd in newbondaries:
-            state = repo.nodestate(nd)
-            if state.trackheads:
-                statesheads.setdefault(state, set([])).add(nd)
-
-        for state, heads in repo._statesheads.iteritems():
-            if not state.trackheads:
-                continue
-            heads = set(heads) - tostrip | statesheads.get(state, set([]))
-            # reduce heads (make them really heads)
-            revs = set(map(cl.rev, heads))
-            minrev = min(revs)
-            for rev in cl.ancestors(*revs):
-                if rev >= minrev:
-                    revs.discard(rev)
-            repo._statesheads[state] = map(cl.node, revs)
-        _writestateshead(repo)
-
-        return ostrip(ui, repo, node, backup)
-
-    mercurial.repair.strip = strip