Discard lines in file diff

keymaps/git.cson

@@ -22,6 +22,7 @@
   'shift-tab': 'github:select-previous-hunk'
   'right': 'core:move-right'
   'o': 'github:open-file'
+  'cmd-z': 'github:undo-last-file-diff-discard'
 
 '.github-Prompt-input':
   'enter': 'core:confirm'

lib/controllers/file-patch-controller.js

@@ -40,18 +40,25 @@ export default class FilePatchController {
       );
     } else {
       // NOTE: Outer div is required for etch to render elements correctly
+      const filePath = this.props.filePatch.getPath();
       return (
         <div className="github-PaneView pane-item">
           <FilePatchView
             ref="filePatchView"
             commandRegistry={this.props.commandRegistry}
             attemptLineStageOperation={this.attemptLineStageOperation}
             didSurfaceFile={this.didSurfaceFile}
+            didDiveIntoCorrespondingFilePatch={this.diveIntoCorrespondingFilePatch}
             attemptHunkStageOperation={this.attemptHunkStageOperation}
             hunks={hunks}
+            filePath={filePath}
             stagingStatus={this.props.stagingStatus}
+            isPartiallyStaged={this.props.isPartiallyStaged}
             registerHunkView={this.props.registerHunkView}
             openCurrentFile={this.openCurrentFile}
+            discardLines={this.props.discardLines}
+            undoLastDiscard={this.undoLastDiscard}
+            hasUndoHistory={this.props.hasUndoHistory}
           />
         </div>
       );
@@ -164,6 +171,13 @@ export default class FilePatchController {
     }
   }
 
+  @autobind
+  diveIntoCorrespondingFilePatch() {
+    const filePath = this.props.filePatch.getPath();
+    const stagingStatus = this.props.stagingStatus === 'staged' ? 'unstaged' : 'staged';
+    return this.props.didDiveIntoFilePath(filePath, stagingStatus, {amending: this.props.isAmending});
+  }
+
   didUpdateFilePatch() {
     // FilePatch was mutated so all we need to do is re-render
     return etch.update(this);
@@ -187,4 +201,9 @@ export default class FilePatchController {
     textEditor.setCursorBufferPosition(position);
     return textEditor;
   }
+
+  @autobind
+  undoLastDiscard() {
+    return this.props.undoLastDiscard(this.props.filePatch.getPath());
+  }
 }

lib/controllers/git-controller.js

@@ -1,6 +1,6 @@
 import path from 'path';
 
-import {CompositeDisposable, Disposable, File} from 'atom';
+import {CompositeDisposable, Disposable, File, TextBuffer} from 'atom';
 
 import React from 'react';
 import {autobind} from 'core-decorators';
@@ -18,11 +18,13 @@ import GitPanelController from './git-panel-controller';
 import StatusBarTileController from './status-bar-tile-controller';
 import ModelObserver from '../models/model-observer';
 import ModelStateRegistry from '../models/model-state-registry';
+import discardChangesInBuffer from '../discard-changes-in-buffer';
 
 const nullFilePatchState = {
   filePath: null,
   filePatch: null,
   stagingStatus: null,
+  partiallyStaged: null,
 };
 
 export default class GitController extends React.Component {
@@ -166,6 +168,7 @@ export default class GitController extends React.Component {
   }
 
   renderFilePatchController() {
+    const hasUndoHistory = this.hasUndoHistory();
     return (
       <div>
         <Commands registry={this.props.commandRegistry} target="atom-workspace">
@@ -182,9 +185,15 @@ export default class GitController extends React.Component {
               commandRegistry={this.props.commandRegistry}
               filePatch={this.state.filePatch}
               stagingStatus={this.state.stagingStatus}
+              isAmending={this.state.amending}
+              isPartiallyStaged={this.state.partiallyStaged}
               onRepoRefresh={this.onRepoRefresh}
               didSurfaceFile={this.surfaceFromFileAtPath}
+              didDiveIntoFilePath={this.diveIntoFilePatchForPath}
               openFiles={this.openFiles}
+              discardLines={this.discardLines}
+              undoLastDiscard={this.undoLastDiscard}
+              hasUndoHistory={hasUndoHistory}
             />
           </EtchWrapper>
         </PaneItem>
@@ -210,9 +219,10 @@ export default class GitController extends React.Component {
     if (!repository) { return null; }
 
     const filePatch = await repository.getFilePatchForPath(filePath, {staged: stagingStatus === 'staged', amending});
+    const partiallyStaged = await repository.isPartiallyStaged(filePath);
     return new Promise(resolve => {
       if (filePatch) {
-        this.setState({filePath, filePatch, stagingStatus}, () => {
+        this.setState({filePath, filePatch, stagingStatus, partiallyStaged}, () => {
           // TODO: can be better done w/ a prop?
           if (activate && this.filePatchControllerPane) {
             this.filePatchControllerPane.activate();
@@ -335,4 +345,70 @@ export default class GitController extends React.Component {
   handleChangeTab(activeTab) {
     this.setState({activeTab});
   }
+
+  @autobind
+  async discardLines(lines) {
+    const relFilePath = this.state.filePatch.getPath();
+    const absfilePath = path.join(this.props.repository.getWorkingDirectoryPath(), relFilePath);
+    let buffer, disposable;
+    const isSafe = async () => {
+      const editor = this.props.workspace.getTextEditors().find(e => e.getPath() === absfilePath);
+      if (editor) {
+        buffer = editor.getBuffer();
+        if (buffer.isModified()) {
+          this.props.notificationManager.addError('Cannot discard lines.', {description: 'You have unsaved changes.'});
+          return false;
+        }
+      } else {
+        buffer = new TextBuffer({filePath: absfilePath, load: true});
+        await new Promise(resolve => {
+          disposable = buffer.onDidReload(() => {
+            disposable.dispose();
+            resolve();
+          });
+        });
+      }
+      return true;
+    };
+    await this.props.repository.storeBeforeAndAfterBlobs(relFilePath, isSafe, () => {
+      this.discardChangesInBuffer(buffer, this.state.filePatch, lines);
+    });
+    if (disposable) { disposable.dispose(); }
+  }
+
+  discardChangesInBuffer(buffer, filePatch, lines) {
+    discardChangesInBuffer(buffer, filePatch, lines);
+  }
+
+  @autobind
+  async undoLastDiscard(filePath) {
+    const relFilePath = this.state.filePatch.getPath();
+    const absfilePath = path.join(this.props.repository.getWorkingDirectoryPath(), relFilePath);
+    const isSafe = () => {
+      const editor = this.props.workspace.getTextEditors().find(e => e.getPath() === absfilePath);
+      if (editor && editor.getBuffer().isModified()) {
+        this.props.notificationManager.addError(
+          'Cannot undo last discard.',
+          {description: 'You have unsaved changes.'},
+        );
+        return false;
+      }
+      return true;
+    };
+    try {
+      await this.props.repository.attemptToRestoreBlob(filePath, isSafe);
+    } catch (e) {
+      if (e.message === 'Cannot restore file. Contents have been modified since last discard.') {
+        this.props.notificationManager.addError(
+          'Cannot undo last discard.',
+          {description: 'Contents have been modified since last discard.'},
+        );
+      }
+    }
+  }
+
+  @autobind
+  hasUndoHistory() {
+    return this.props.repository.hasUndoHistory(this.state.filePatch.getPath());
+  }
 }

lib/discard-changes-in-buffer.js

@@ -0,0 +1,35 @@
+export default function discardChangesInBuffer(buffer, filePatch, discardedLines) {
+  buffer.transact(() => {
+    let addedCount = 0;
+    let removedCount = 0;
+    let deletedCount = 0;
+    filePatch.getHunks().forEach(hunk => {
+      hunk.getLines().forEach(line => {
+        if (discardedLines.has(line)) {
+          if (line.status === 'deleted') {
+            const row = (line.oldLineNumber - deletedCount) + addedCount - removedCount - 1;
+            buffer.insert([row, 0], line.text + '\n');
+            addedCount++;
+          } else if (line.status === 'added') {
+            const row = line.newLineNumber + addedCount - removedCount - 1;
+            if (buffer.lineForRow(row) === line.text) {
+              buffer.deleteRow(row);
+              removedCount++;
+            } else {
+              throw new Error(buffer.lineForRow(row) + ' does not match ' + line.text);
+            }
+          } else if (line.status === 'nonewline') {
+            // TODO: handle no new line case
+          } else {
+            throw new Error(`unrecognized status: ${line.status}. Must be 'added' or 'deleted'`);
+          }
+        }
+        if (line.getStatus() === 'deleted') {
+          deletedCount++;
+        }
+      });
+    });
+  });
+
+  buffer.save();
+}

lib/git-shell-out-strategy.js

@@ -7,7 +7,7 @@ import {parse as parseDiff} from 'what-the-diff';
 
 import GitPromptServer from './git-prompt-server';
 import AsyncQueue from './async-queue';
-import {readFile, fsStat, deleteFileOrFolder} from './helpers';
+import {readFile, writeFile, fsStat, deleteFileOrFolder} from './helpers';
 
 const LINE_ENDING_REGEX = /\r?\n/;
 
@@ -36,7 +36,7 @@ export default class GitShellOutStrategy {
   }
 
   // Execute a command and read the output using the embedded Git environment
-  exec(args, stdin = null, useGitPromptServer = false) {
+  exec(args, stdin = null, useGitPromptServer = false, redirectFilePath = null) {
     /* eslint-disable no-console */
     const subscriptions = new CompositeDisposable();
     return this.commandQueue.push(async () => {
@@ -99,6 +99,9 @@ export default class GitShellOutStrategy {
             err.command = formattedArgs;
             return Promise.reject(err);
           }
+          if (redirectFilePath) {
+            return writeFile(redirectFilePath, stdout);
+          }
           return stdout;
         });
     });
@@ -268,6 +271,19 @@ export default class GitShellOutStrategy {
     return rawDiffs[0];
   }
 
+  async isPartiallyStaged(filePath) {
+    const args = ['status', '--short', '--', filePath];
+    const output = await this.exec(args);
+    const results = output.trim().split(LINE_ENDING_REGEX);
+    if (results.length === 2) {
+      return true;
+    } else if (results.length === 1) {
+      return ['MM', 'AM', 'MD'].includes(results[0].slice(0, 2));
+    } else {
+      throw new Error(`Unexpected output for ${args.join(' ')}: ${output}`);
+    }
+  }
+
   /**
    * Miscellaneous getters
    */
@@ -445,6 +461,15 @@ export default class GitShellOutStrategy {
       return [];
     }
   }
+
+  async createBlob(filePath) {
+    const output = await this.exec(['hash-object', '-w', filePath]);
+    return output.trim();
+  }
+
+  async restoreBlob(filePath, sha) {
+    await this.exec(['cat-file', '-p', sha], null, null, path.join(this.workingDir, filePath));
+  }
 }
 
 function buildAddedFilePatch(filePath, contents, stats) {

lib/helpers.js

@@ -9,6 +9,14 @@ export function readFile(absoluteFilePath, encoding = 'utf8') {
   });
 }
 
+export function writeFile(absoluteFilePath, contents) {
+  return new Promise((resolve, reject) => {
+    fs.writeFile(absoluteFilePath, contents, err => {
+      if (err) { return reject(err); } else { return resolve(); }
+    });
+  });
+}
+
 export function deleteFileOrFolder(path) {
   return new Promise((resolve, reject) => {
     fs.remove(path, err => {

lib/models/file-discard-history.js

@@ -0,0 +1,39 @@
+export default class FileDiscardHistory {
+  constructor(createBlob, restoreBlob) {
+    this.createBlob = createBlob;
+    this.restoreBlob = restoreBlob;
+    this.blobHistoryByFilePath = new Map();
+  }
+
+  async storeBlobs(filePath, isSafe, destructiveAction) {
+    const beforeSha = await this.createBlob(filePath);
+    const isNotSafe = !(await isSafe());
+    if (isNotSafe) { return; }
+    destructiveAction();
+    const afterSha = await this.createBlob(filePath);
+    const snapshots = {beforeSha, afterSha};
+    const history = this.blobHistoryByFilePath.get(filePath);
+    if (history) {
+      history.push(snapshots);
+    } else {
+      this.blobHistoryByFilePath.set(filePath, [snapshots]);
+    }
+  }
+
+  async attemptToRestoreBlob(filePath, isSafe) {
+    const history = this.blobHistoryByFilePath.get(filePath);
+    const {beforeSha, afterSha} = history[history.length - 1];
+    const currentSha = await this.createBlob(filePath);
+    if (currentSha === afterSha && isSafe()) {
+      await this.restoreBlob(filePath, beforeSha);
+      history.pop();
+    } else {
+      throw new Error('Cannot restore file. Contents have been modified since last discard.');
+    }
+  }
+
+  hasUndoHistory(filePath) {
+    const history = this.blobHistoryByFilePath.get(filePath);
+    return !!history && history.length > 0;
+  }
+}