Merge pull request #302 from denisbondar/main

Remove all ANSI codes from buf
carlmontanari · Oct 21, 2023 · d0eb38b · d0eb38b
2 parents b67ea75 + 7358489
commit d0eb38b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/scrapli/channel/base_channel.py b/scrapli/channel/base_channel.py
@@ -10,7 +10,15 @@
 from scrapli.logging import get_instance_logger
 from scrapli.transport.base import AsyncTransport, Transport
 
-ANSI_ESCAPE_PATTERN = re.compile(rb"\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\)|E)")
+ANSI_ESCAPE_PATTERN = re.compile(
+    rb"[\x1B\x9B]"
+    rb"[\[\]()#;?]*"
+    rb"("
+    rb"(([a-zA-Z0-9]*(;[a-zA-Z\d]*)*)?\x07)"
+    rb"|"
+    rb"((\d{1,4}(;\d{0,4})*)?[\dA-PRZcf-ntqry=><~])"
+    rb")"
+)
 
 
 @dataclass()

diff --git a/tests/unit/channel/test_base_channel.py b/tests/unit/channel/test_base_channel.py
@@ -292,6 +292,14 @@ def test_process_output(base_channel):
             b"VeryLong\x1bECommand",
             b"VeryLongCommand",
         ),
+        (
+            b"command-prompt# \x1b7",
+            b"command-prompt# ",
+        ),
+        (
+            b"\x1b[7mCTRL+C\x1b[0m \x1b[7mESC\x1b[0m \x1b[7mq\x1b[0m Quit \x1b[7mSPACE\x1b[0m \x1b[7mn\x1b[0m Next Page \x1b[7mENTER\x1b[0m Next Entry \x1b[7ma\x1b[0m All\x1b[1A\x1b[59C\x1b[27m",
+            b"CTRL+C ESC q Quit SPACE n Next Page ENTER Next Entry a All",
+        ),
     ),
 )
 def test_strip_ansi(base_channel, buf: bytes, expected: bytes):