From 180ff20be59e15f5df3289cf7c3d6d2a29a1c8ff Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Mon, 6 Jan 2020 09:23:53 -0800
Subject: [PATCH] Add :sort command

---
 babi.py                    | 28 ++++++++++++++++++++++++
 testing/runner.py          | 11 ++++++++++
 tests/command_mode_test.py | 12 +---------
 tests/sort_test.py         | 45 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 85 insertions(+), 11 deletions(-)
 create mode 100644 tests/sort_test.py

diff --git a/babi.py b/babi.py
index ef493cb..3213481 100644
--- a/babi.py
+++ b/babi.py
@@ -6,6 +6,7 @@ import enum
 import functools
 import hashlib
 import io
+import itertools
 import os
 import re
 import signal
@@ -1051,6 +1052,27 @@ class File:
         self.x = self.x_hint = len(self.lines[self.y])
         self.lines[self.y] += self.lines.pop(self.y + 1)
 
+    def _sort(self, margin: Margin, s_y: int, e_y: int) -> None:
+        # self.lines intentionally does not support slicing so we use islice
+        lines = sorted(itertools.islice(self.lines, s_y, e_y))
+        for i, line in zip(range(s_y, e_y), lines):
+            self.lines[i] = line
+
+        self.y = s_y
+        self.x = self.x_hint = 0
+        self.scroll_screen_if_needed(margin)
+
+    @edit_action('sort', final=True)
+    def sort(self, margin: Margin) -> None:
+        self._sort(margin, 0, len(self.lines) - 1)
+
+    @edit_action('sort selection', final=True)
+    @clear_selection
+    def sort_selection(self, margin: Margin) -> None:
+        (s_y, _), (e_y, _) = self._get_selection()
+        e_y = min(e_y + 1, len(self.lines) - 1)
+        self._sort(margin, s_y, e_y)
+
     DISPATCH = {
         # movement
         b'KEY_UP': up,
@@ -1448,6 +1470,12 @@ class Screen:
         elif response == ':wq':
             self.save()
             return EditResult.EXIT
+        elif response == ':sort':
+            if self.file.select_start:
+                self.file.sort_selection(self.margin)
+            else:
+                self.file.sort(self.margin)
+            self.status.update('sorted!')
         elif response is not PromptResult.CANCELLED:
             self.status.update(f'invalid command: {response}')
         return None
diff --git a/testing/runner.py b/testing/runner.py
index c00f637..6c77f90 100644
--- a/testing/runner.py
+++ b/testing/runner.py
@@ -143,3 +143,14 @@ def and_exit(h):
     if ' *' in h.get_screen_line(0):
         h.press('n')
     h.await_exit()
+
+
+def trigger_command_mode(h):
+    # in order to enter a steady state, trigger an unknown key first and then
+    # press escape to open the command mode.  this is necessary as `Escape` is
+    # the start of "escape sequences" and sending characters too quickly will
+    # be interpreted as a single keypress
+    h.press('^J')
+    h.await_text('unknown key')
+    h.press('Escape')
+    h.await_text_missing('unknown key')
diff --git a/tests/command_mode_test.py b/tests/command_mode_test.py
index 7fe31d6..16e30fc 100644
--- a/tests/command_mode_test.py
+++ b/tests/command_mode_test.py
@@ -2,17 +2,7 @@ import pytest
 
 from testing.runner import and_exit
 from testing.runner import run
-
-
-def trigger_command_mode(h):
-    # in order to enter a steady state, trigger an unknown key first and then
-    # press escape to open the command mode.  this is necessary as `Escape` is
-    # the start of "escape sequences" and sending characters too quickly will
-    # be interpreted as a single keypress
-    h.press('^J')
-    h.await_text('unknown key')
-    h.press('Escape')
-    h.await_text_missing('unknown key')
+from testing.runner import trigger_command_mode
 
 
 def test_quit_via_colon_q():
diff --git a/tests/sort_test.py b/tests/sort_test.py
new file mode 100644
index 0000000..80912f6
--- /dev/null
+++ b/tests/sort_test.py
@@ -0,0 +1,45 @@
+import pytest
+
+from testing.runner import and_exit
+from testing.runner import run
+from testing.runner import trigger_command_mode
+
+
+@pytest.fixture
+def unsorted(tmpdir):
+    f = tmpdir.join('f')
+    f.write('d\nb\nc\na\n')
+    return f
+
+
+def test_sort_entire_file(unsorted):
+    with run(str(unsorted)) as h, and_exit(h):
+        trigger_command_mode(h)
+        h.press_and_enter(':sort')
+        h.await_text('sorted!')
+        h.await_cursor_position(x=0, y=1)
+        h.press('^S')
+    assert unsorted.read() == 'a\nb\nc\nd\n'
+
+
+def test_sort_selection(unsorted):
+    with run(str(unsorted)) as h, and_exit(h):
+        h.press('S-Down')
+        trigger_command_mode(h)
+        h.press_and_enter(':sort')
+        h.await_text('sorted!')
+        h.await_cursor_position(x=0, y=1)
+        h.press('^S')
+    assert unsorted.read() == 'b\nd\nc\na\n'
+
+
+def test_sort_selection_does_not_include_eof(unsorted):
+    with run(str(unsorted)) as h, and_exit(h):
+        for _ in range(5):
+            h.press('S-Down')
+        trigger_command_mode(h)
+        h.press_and_enter(':sort')
+        h.await_text('sorted!')
+        h.await_cursor_position(x=0, y=1)
+        h.press('^S')
+    assert unsorted.read() == 'a\nb\nc\nd\n'