1 files changed, 134 insertions, 42 deletions

diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/api/CheckoutCommand.java b/org.eclipse.jgit/src/org/eclipse/jgit/api/CheckoutCommand.java
index 9d3bf561b5..ef227dc854 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/api/CheckoutCommand.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/api/CheckoutCommand.java
@@ -125,6 +125,33 @@ import org.eclipse.jgit.util.FileUtils;
  *      >Git documentation about Checkout</a>
  */
 public class CheckoutCommand extends GitCommand<Ref> {
+
+	/**
+	 * Stage to check out, see {@link CheckoutCommand#setStage(Stage)}.
+	 */
+	public static enum Stage {
+		/**
+		 * Base stage (#1)
+		 */
+		BASE(DirCacheEntry.STAGE_1),
+
+		/**
+		 * Ours stage (#2)
+		 */
+		OURS(DirCacheEntry.STAGE_2),
+
+		/**
+		 * Theirs stage (#3)
+		 */
+		THEIRS(DirCacheEntry.STAGE_3);
+
+		private final int number;
+
+		private Stage(int number) {
+			this.number = number;
+		}
+	}
+
 	private String name;
 
 	private boolean force = false;
@@ -137,6 +164,8 @@ public class CheckoutCommand extends GitCommand<Ref> {
 
 	private RevCommit startCommit;
 
+	private Stage checkoutStage = null;
+
 	private CheckoutResult status;
 
 	private List<String> paths;
@@ -327,52 +356,19 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		RevWalk revWalk = new RevWalk(repo);
 		DirCache dc = repo.lockDirCache();
 		try {
-			DirCacheEditor editor = dc.editor();
-			TreeWalk startWalk = new TreeWalk(revWalk.getObjectReader());
-			startWalk.setRecursive(true);
+			TreeWalk treeWalk = new TreeWalk(revWalk.getObjectReader());
+			treeWalk.setRecursive(true);
 			if (!checkoutAllPaths)
-				startWalk.setFilter(PathFilterGroup.createFromStrings(paths));
-			final boolean checkoutIndex = startCommit == null
-					&& startPoint == null;
-			if (!checkoutIndex)
-				startWalk.addTree(revWalk.parseCommit(getStartPoint())
-						.getTree());
-			else
-				startWalk.addTree(new DirCacheIterator(dc));
-
-			final File workTree = repo.getWorkTree();
-			final ObjectReader r = repo.getObjectDatabase().newReader();
+				treeWalk.setFilter(PathFilterGroup.createFromStrings(paths));
 			try {
-				while (startWalk.next()) {
-					final ObjectId blobId = startWalk.getObjectId(0);
-					final FileMode mode = startWalk.getFileMode(0);
-					editor.add(new PathEdit(startWalk.getPathString()) {
-						public void apply(DirCacheEntry ent) {
-							if (checkoutIndex
-									&& ent.getStage() > DirCacheEntry.STAGE_0) {
-								UnmergedPathException e = new UnmergedPathException(ent);
-								throw new JGitInternalException(e.getMessage(), e);
-							}
-							ent.setObjectId(blobId);
-							ent.setFileMode(mode);
-							File file = new File(workTree, ent.getPathString());
-							File parentDir = file.getParentFile();
-							try {
-								FileUtils.mkdirs(parentDir, true);
-								DirCacheCheckout.checkoutEntry(repo, file, ent, r);
-							} catch (IOException e) {
-								throw new JGitInternalException(
-										MessageFormat.format(
-												JGitText.get().checkoutConflictWithFile,
-												ent.getPathString()), e);
-							}
-						}
-					});
+				if (isCheckoutIndex())
+					checkoutPathsFromIndex(treeWalk, dc);
+				else {
+					RevCommit commit = revWalk.parseCommit(getStartPoint());
+					checkoutPathsFromCommit(treeWalk, dc, commit);
 				}
-				editor.commit();
 			} finally {
-				startWalk.release();
-				r.release();
+				treeWalk.release();
 			}
 		} finally {
 			dc.unlock();
@@ -381,6 +377,75 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		return this;
 	}
 
+	private void checkoutPathsFromIndex(TreeWalk treeWalk, DirCache dc)
+			throws IOException {
+		DirCacheIterator dci = new DirCacheIterator(dc);
+		treeWalk.addTree(dci);
+
+		final ObjectReader r = treeWalk.getObjectReader();
+		DirCacheEditor editor = dc.editor();
+		while (treeWalk.next()) {
+			DirCacheEntry entry = dci.getDirCacheEntry();
+			// Only add one edit per path
+			if (entry != null && entry.getStage() > DirCacheEntry.STAGE_1)
+				continue;
+			editor.add(new PathEdit(treeWalk.getPathString()) {
+				public void apply(DirCacheEntry ent) {
+					int stage = ent.getStage();
+					if (stage > DirCacheEntry.STAGE_0) {
+						if (checkoutStage != null) {
+							if (stage == checkoutStage.number)
+								checkoutPath(ent, r);
+						} else {
+							UnmergedPathException e = new UnmergedPathException(
+									ent);
+							throw new JGitInternalException(e.getMessage(), e);
+						}
+					} else {
+						checkoutPath(ent, r);
+					}
+				}
+			});
+		}
+		editor.commit();
+	}
+
+	private void checkoutPathsFromCommit(TreeWalk treeWalk, DirCache dc,
+			RevCommit commit) throws IOException {
+		treeWalk.addTree(commit.getTree());
+		final ObjectReader r = treeWalk.getObjectReader();
+		DirCacheEditor editor = dc.editor();
+		while (treeWalk.next()) {
+			final ObjectId blobId = treeWalk.getObjectId(0);
+			final FileMode mode = treeWalk.getFileMode(0);
+			editor.add(new PathEdit(treeWalk.getPathString()) {
+				public void apply(DirCacheEntry ent) {
+					ent.setObjectId(blobId);
+					ent.setFileMode(mode);
+					checkoutPath(ent, r);
+				}
+			});
+		}
+		editor.commit();
+	}
+
+	private void checkoutPath(DirCacheEntry entry, ObjectReader reader) {
+		File file = new File(repo.getWorkTree(), entry.getPathString());
+		File parentDir = file.getParentFile();
+		try {
+			FileUtils.mkdirs(parentDir, true);
+			DirCacheCheckout.checkoutEntry(repo, file, entry, reader);
+		} catch (IOException e) {
+			throw new JGitInternalException(MessageFormat.format(
+					JGitText.get().checkoutConflictWithFile,
+					entry.getPathString()), e);
+		}
+	}
+
+	private boolean isCheckoutIndex() {
+		return startCommit == null && startPoint == null;
+	}
+
 	private ObjectId getStartPoint() throws AmbiguousObjectException,
 			RefNotFoundException, IOException {
 		if (startCommit != null)
@@ -483,6 +548,7 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		checkCallable();
 		this.startPoint = startPoint;
 		this.startCommit = null;
+		checkOptions();
 		return this;
 	}
 
@@ -503,6 +569,7 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		checkCallable();
 		this.startCommit = startCommit;
 		this.startPoint = null;
+		checkOptions();
 		return this;
 	}
 
@@ -523,6 +590,24 @@ public class CheckoutCommand extends GitCommand<Ref> {
 	}
 
 	/**
+	 * When checking out the index, check out the specified stage (ours or
+	 * theirs) for unmerged paths.
+	 * <p>
+	 * This can not be used when checking out a branch, only when checking out
+	 * the index.
+	 *
+	 * @param stage
+	 *            the stage to check out
+	 * @return this
+	 */
+	public CheckoutCommand setStage(Stage stage) {
+		checkCallable();
+		this.checkoutStage = stage;
+		checkOptions();
+		return this;
+	}
+
+	/**
 	 * @return the result, never <code>null</code>
 	 */
 	public CheckoutResult getResult() {
@@ -530,4 +615,11 @@ public class CheckoutCommand extends GitCommand<Ref> {
 			return CheckoutResult.NOT_TRIED_RESULT;
 		return status;
 	}
+
+	private void checkOptions() {
+		if (checkoutStage != null && !isCheckoutIndex())
+			throw new IllegalStateException(
+					"Checking out ours/theirs is only possible when checking out index, "
+							+ "not when switching branches.");
+	}
 }