snapshot: add --clean option to snapshot cmd

Summary:
It cleans the working copy, equivalent of `hg checkout --clean && hg
purge`.

Reviewed By: markbt

Differential Revision: D16831349

fbshipit-source-id: d765f2f36f561786ec57bc1e54fdf81695f1e7e0

edenscm/hgext/snapshot.py

@@ -30,6 +30,7 @@ from edenscm.mercurial import (
     pathutil,
     registrar,
     scmutil,
+    util,
     visibility,
 )
 from edenscm.mercurial.i18n import _
@@ -191,12 +192,30 @@ class snapshotmanifest(object):
         lfs.wrapper.uploadblobs(self.repo, pointers)
 
 
-@command("debugsnapshot", inferrepo=True)
+@command(
+    "debugsnapshot", [("", "clean", False, _("clean the working copy"))], inferrepo=True
+)
 def debugsnapshot(ui, repo, *args, **opts):
     """
     Creates a snapshot of the working copy.
     TODO(alexeyqu): finish docs
     """
+
+    def removeuntrackedfiles(ui, repo):
+        """
+        Removes all untracked files from the repo.
+        """
+        # the same behavior is implemented better in the purge extension
+        # more corner cases are handled there
+        # e.g. directories that became empty during purge get deleted too
+        # TODO(alexeyqu): use code from purge, probable move it to core code
+        status = repo.status(unknown=True)
+        for file in status.unknown:
+            try:
+                util.tryunlink(repo.wjoin(file))
+            except OSError:
+                ui.warn(_("%s cannot be removed") % file)
+
     with repo.wlock(), repo.lock():
         node = createsnapshotcommit(ui, repo, opts)
         if not node:
@@ -205,6 +224,14 @@ def debugsnapshot(ui, repo, *args, **opts):
         ui.status(_("snapshot %s created\n") % (repo[node].hex()))
         if visibility.enabled(repo):
             visibility.remove(repo, [node])
+        if opts.get("clean"):
+            try:
+                # We want to bring the working copy to the p1 state
+                rev = repo[None].p1()
+                hg.updatetotally(ui, repo, rev, rev, clean=True)
+                removeuntrackedfiles(ui, repo)
+            except (KeyboardInterrupt, Exception) as exc:
+                ui.warn(_("failed to clean the working copy: %s\n") % exc)
 
 
 def createsnapshotcommit(ui, repo, opts):

tests/test-fb-hgext-snapshot.t

@@ -31,12 +31,13 @@
 
 # Snapshot with empty manifest (changes only in tracked files)
   $ hg rm bar/file
-  $ EMPTYOID="$(hg debugsnapshot | cut -f2 -d' ')"
+  $ EMPTYOID="$(hg debugsnapshot --clean | head -n 1 | cut -f2 -d' ')"
+  $ echo "$EMPTYOID"
+  669ab2753dcee94dbdb1aec45db328db0f1d81f5
   $ hg log --hidden -r "$EMPTYOID" -T '{extras % \"{extra}\n\"}' | grep snapshotmanifestid
   snapshotmanifestid=None
 
 # Merge conflict!
-  $ hg revert bar/file
   $ echo "a" > mergefile
   $ hg add mergefile
   $ hg commit -m "merge #1"
@@ -193,6 +194,24 @@
   # To continue:                hg commit
   # To abort:                   hg update --clean .    (warning: this will discard uncommitted changes)
   
+
+# Oh wait, we need to clean all that for some reason
+  $ find .hg/merge | sort
+  .hg/merge
+  .hg/merge/fc4ffdcb8ed23cecd44a0e11d23af83b445179b4
+  .hg/merge/state
+  .hg/merge/state2
+  $ hg debugsnapshot --clean
+  snapshot e203bc3b0e5140c52ec5ba01fcb7071af5ea40de created
+  3 files updated, 0 files merged, 0 files removed, 0 files unresolved
+  $ hg status --verbose
+  $ test -d .hg/merge
+  [1]
+  $ hg debugcheckoutsnapshot --hidden "$OID"
+  will checkout on ccdff83036b6b05c657a1eebff7dc523b865f6ce
+  2 files updated, 0 files merged, 0 files removed, 0 files unresolved
+  checkout complete
+
 # Finally, resolve the conflict
   $ hg resolve --mark mergefile
   (no more unresolved files)