From 4921d0fea258fd86d38da497a636ac39b62fd818 Mon Sep 17 00:00:00 2001
From: Daniel Miller <millerdev@gmail.com>
Date: Sun, 7 Feb 2021 13:58:03 -0500
Subject: [PATCH] Fix ag command handling of very long line matches

---
 CHANGELOG.md               |  5 +++++
 package.json               |  2 +-
 pyxt/cmd/ag.py             |  6 +++++-
 pyxt/cmd/tests/test_ag.py  |  6 +++---
 pyxt/process.py            | 36 ++++++++++++++++++++++++++++++++++--
 pyxt/tests/test_process.py | 15 +++++++++++++++
 6 files changed, 63 insertions(+), 7 deletions(-)
 create mode 100644 pyxt/tests/test_process.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7a14eca..3c00a0b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,8 @@
+0.3.3 - tbd
+
+- Fix `ag` command handling of very long line matches.
+
+
 0.3.2 - 2021-02-05
 
 - Fix isort entire document (not selection).
diff --git a/package.json b/package.json
index c71a370..1112f73 100644
--- a/package.json
+++ b/package.json
@@ -6,7 +6,7 @@
   "repository": "https://github.com/millerdev/pyxt",
   "license": "See LICENSE file",
   "homepage": "https://github.com/millerdev/pyxt",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "publisher": "millerdev",
   "engines": {
     "vscode": "^1.34.0"
diff --git a/pyxt/cmd/ag.py b/pyxt/cmd/ag.py
index 7f61f5e..19fbc6e 100644
--- a/pyxt/cmd/ag.py
+++ b/pyxt/cmd/ag.py
@@ -118,7 +118,11 @@ def got_output(item, returncode, error=""):
         if len(items) >= MAX_RESULT_ITEMS:
             raise TooManyResults
 
-    return {"iter_output": ag_lines, "got_output": got_output}
+    return {
+        "iter_output": ag_lines,
+        "got_output": got_output,
+        "limit": MAX_LINE_LENGTH,
+    }
 
 
 def create_item(abspath, relpath, num, ranges, delim, text):
diff --git a/pyxt/cmd/tests/test_ag.py b/pyxt/cmd/tests/test_ag.py
index e026f30..a6751c7 100644
--- a/pyxt/cmd/tests/test_ag.py
+++ b/pyxt/cmd/tests/test_ag.py
@@ -117,11 +117,11 @@ async def test_huge_results():
     with tempdir() as tmp:
         os.mkdir(join(tmp, "dir"))
         with open(join(tmp, "dir/long_lines.txt"), "w", encoding="utf-8") as fh:
-            fh.write("file.txt " * 100)
+            fh.write("xyz " + "file.txt " * 100)
         editor = FakeEditor(join(tmp, "dir/file"), tmp)
-        result = await do_command("ag txt", editor)
+        result = await do_command("ag xyz", editor)
         items = result["items"]
-        eq(len(items[0]["label"]), mod.MAX_LINE_LENGTH + 3, result)
+        eq(len(items[0]["label"]), mod.MAX_LINE_LENGTH - 3, result)
         eq(len(items), 1, result)
 
 
diff --git a/pyxt/process.py b/pyxt/process.py
index d451853..2437c84 100644
--- a/pyxt/process.py
+++ b/pyxt/process.py
@@ -1,5 +1,5 @@
 import logging
-from asyncio.exceptions import CancelledError
+from asyncio.exceptions import CancelledError, IncompleteReadError, LimitOverrunError
 from asyncio.subprocess import create_subprocess_exec, PIPE, STDOUT
 
 log = logging.getLogger(__name__)
@@ -68,11 +68,43 @@ def got_output(line, returncode, error=""):
 
 async def iter_lines(stream, encoding):
     while True:
-        line = await stream.readline()
+        line = await readline(stream)
         if not line:
             break
         yield line.decode(encoding)
 
 
+async def readline(stream):
+    # adapted from StreamReader.readline
+    sep = b'\n'
+    try:
+        line = await stream.readuntil(sep)
+    except IncompleteReadError as e:
+        return e.partial
+    except LimitOverrunError as e:
+        line = stream._buffer[:stream._limit]
+        await _discard_overrun(stream, sep, e.consumed)
+    return line
+
+
+async def _discard_overrun(stream, sep, consumed):
+    seplen = len(sep)
+    while True:
+        if stream._buffer.startswith(sep, consumed):
+            del stream._buffer[:consumed + seplen]
+            found_sep = True
+        else:
+            stream._buffer.clear()
+            found_sep = False
+        stream._maybe_resume_transport()
+        if not found_sep:
+            try:
+                await stream.readuntil(sep)
+            except LimitOverrunError as e:
+                consumed = e.consumed
+                continue
+        break
+
+
 class ProcessError(Exception):
     pass
diff --git a/pyxt/tests/test_process.py b/pyxt/tests/test_process.py
new file mode 100644
index 0000000..bdc6a20
--- /dev/null
+++ b/pyxt/tests/test_process.py
@@ -0,0 +1,15 @@
+from testil import eq
+
+from .util import async_test
+from ..process import process_lines
+
+
+@async_test
+async def test_process_lines_length_limit():
+    def got_output(line, code):
+        if line is not None:
+            lines.append(line)
+    lines = []
+    cmd = ["echo", "line 1\nline 1 000 000 000\nline 20\nline 30"]
+    await process_lines(cmd, got_output=got_output, limit=6)
+    eq(lines, ["line 1\n", "line 1", "line 2", "line 3"])