Skip to content

Commit ff66353

Browse files
committed
Added strip_ansi() function and visible_prompt propery
Transcript testing now strips ANSI escape codes from the current line and current prompt before comparing the start of the current line with the current prompt to see if it is a command. This makes it easier to use transcript testing with applications that colorize the prompt.
1 parent 9e74cfd commit ff66353

File tree

1 file changed

+30
-6
lines changed

1 file changed

+30
-6
lines changed

cmd2.py

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,19 @@ class EmptyStatement(Exception):
389389
pass
390390

391391

392+
# Regular expression to match ANSI escape codes
393+
ANSI_ESCAPE_RE = re.compile(r'\x1b[^m]*m')
394+
395+
396+
def strip_ansi(text):
397+
"""Strip ANSI escape codes from a string.
398+
399+
:param text: str - a string which may contain ANSI escape codes
400+
:return: str - the same string with any ANSI escape codes removed
401+
"""
402+
return ANSI_ESCAPE_RE.sub('', text)
403+
404+
392405
class Cmd(cmd.Cmd):
393406
"""An easy but powerful framework for writing line-oriented command interpreters.
394407
@@ -520,6 +533,17 @@ def __init__(self, completekey='tab', stdin=None, stdout=None, use_ipython=False
520533

521534
# ----- Methods related to presenting output to the user -----
522535

536+
@property
537+
def visible_prompt(self):
538+
"""Read-only property to get the visible prompt with any ANSI escape codes stripped.
539+
540+
Used by transcript testing to make it easier and more reliable when users are doing things like coloring the
541+
prompt using ANSI color codes.
542+
543+
:return: str - prompt stripped of any ANSI escape codes
544+
"""
545+
return strip_ansi(self.prompt)
546+
523547
def _finalize_app_parameters(self):
524548
self.commentGrammars.ignore(pyparsing.quotedString).setParseAction(lambda x: '')
525549
# noinspection PyUnresolvedReferences
@@ -2183,18 +2207,18 @@ def runTest(self): # was testall
21832207
def _test_transcript(self, fname, transcript):
21842208
line_num = 0
21852209
finished = False
2186-
line = next(transcript)
2210+
line = strip_ansi(next(transcript))
21872211
line_num += 1
21882212
while not finished:
21892213
# Scroll forward to where actual commands begin
2190-
while not line.startswith(self.cmdapp.prompt):
2214+
while not line.startswith(self.cmdapp.visible_prompt):
21912215
try:
2192-
line = next(transcript)
2216+
line = strip_ansi(next(transcript))
21932217
except StopIteration:
21942218
finished = True
21952219
break
21962220
line_num += 1
2197-
command = [line[len(self.cmdapp.prompt):]]
2221+
command = [line[len(self.cmdapp.visible_prompt):]]
21982222
line = next(transcript)
21992223
# Read the entirety of a multi-line command
22002224
while line.startswith(self.cmdapp.continuation_prompt):
@@ -2213,13 +2237,13 @@ def _test_transcript(self, fname, transcript):
22132237
self.cmdapp.onecmd_plus_hooks(command)
22142238
result = self.cmdapp.stdout.read()
22152239
# Read the expected result from transcript
2216-
if line.startswith(self.cmdapp.prompt):
2240+
if strip_ansi(line).startswith(self.cmdapp.visible_prompt):
22172241
message = '\nFile %s, line %d\nCommand was:\n%r\nExpected: (nothing)\nGot:\n%r\n' % \
22182242
(fname, line_num, command, result)
22192243
self.assert_(not (result.strip()), message)
22202244
continue
22212245
expected = []
2222-
while not line.startswith(self.cmdapp.prompt):
2246+
while not strip_ansi(line).startswith(self.cmdapp.visible_prompt):
22232247
expected.append(line)
22242248
try:
22252249
line = next(transcript)

0 commit comments

Comments
 (0)