Skip to content

Commit

Permalink
trim_markdown: handle the case when there are no code blocks in the i…
Browse files Browse the repository at this point in the history
…nput; more unit tests
  • Loading branch information
Robert Szefler committed Jul 8, 2024
1 parent 9f054ec commit c721262
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 5 deletions.
22 changes: 18 additions & 4 deletions src/robusta/utils/trim_markdown.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,21 +23,35 @@ def trim_markdown(text: str, max_length: int, suffix: str = "...") -> str:
if trim_idx >= close_end:
# Trimming point after this block quote
return text[:trim_idx] + suffix
if trim_idx < open_start:
elif trim_idx < open_start:
# Trimming point before this block quote - continue to the preceding block
continue
if trim_idx >= open_start and trim_idx < open_start + 3:
elif trim_idx >= open_start and trim_idx < open_start + 3:
# Trimming point inside the opening block quote tag
return text[:trim_idx].rstrip("`") + suffix
if trim_idx >= close_start and trim_idx < close_end:
elif trim_idx >= close_start and trim_idx < close_end:
# Trimming point inside the closing block quote tag
if trim_idx - open_end >= 3: # Enough space to insert the closing tag
return text[:trim_idx - 3] + "```" + suffix
else: # Not enough space, strip the whole block
return text[:open_start] + suffix
if trim_idx >= open_end and trim_idx < close_start:
elif trim_idx >= open_end and trim_idx < close_start:
# Trimming point inside the block quote
if trim_idx - open_end >= 3: # Enough space to insert the closing tag
return text[:trim_idx - 3] + "```" + suffix
else: # Not enough space, strip the whole block
return text[:open_start] + suffix
else:
# This should never happen
raise Exception(
f'Internal error in trim_markdown, text="{text[:12]}"(...), {max_length=}, suffix="{suffix}", '
f'matched code block {open_start}..{close_end}'
)

# Cases when there were no code blocks in the input
if len(text) <= trim_idx:
return text
elif len(text) < max_length:
return (text[:trim_idx] + suffix)[:max_length]
else:
return text[:trim_idx] + suffix
26 changes: 25 additions & 1 deletion tests/test_trim_markdown.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,30 @@
])
def test_trim_markdown(max_length: int, expected_output: str):
text = "```oh``` hello ```world``` and then ```something```"
trimmed = trim_markdown(text, max_length, '##')
trimmed = trim_markdown(text, max_length, "##")
assert trimmed == expected_output
assert len(trimmed) <= max_length


@pytest.mark.parametrize(
"max_length,expected_output", [
(0, ""),
(1, "$"),
(2, "$$"),
(3, "$$$"),
(4, "N$$$"),
(5, "No$$$"),
(10, "No code$$$"),
(38, "No code blocks whatsoever in this t$$$"),
(39, "No code blocks whatsoever in this te$$$"),
(40, "No code blocks whatsoever in this tex$$$"),
(41, "No code blocks whatsoever in this text"),
(42, "No code blocks whatsoever in this text"),
(111, "No code blocks whatsoever in this text"),
]
)
def test_trim_markdown_no_code_blocks(max_length: int, expected_output: str):
text = "No code blocks whatsoever in this text"
trimmed = trim_markdown(text, max_length, "$$$")
assert trimmed == expected_output
assert len(trimmed) <= max_length

0 comments on commit c721262

Please sign in to comment.