Create a utility function to count lines in a string

3 files changed, 13 insertions, 2 deletions

diff --git a/bot/cogs/codeblock/cog.py b/bot/cogs/codeblock/cog.py
index 74f122936..ecaf51aa0 100644
--- a/bot/cogs/codeblock/cog.py
+++ b/bot/cogs/codeblock/cog.py
@@ -8,6 +8,7 @@ from discord.ext.commands import Bot, Cog
 
 from bot.cogs.token_remover import TokenRemover
 from bot.constants import Channels, DEBUG_MODE
+from bot.utils import has_lines
 from bot.utils.channel import is_help_channel
 from bot.utils.messages import wait_for_deletion
 from .instructions import get_instructions
@@ -134,7 +135,7 @@ class CodeBlockCog(Cog, name="Code Block"):
         return (
             not message.author.bot
             and self.is_valid_channel(message.channel)
-            and len(message.content.split("\n", 3)) > 3
+            and has_lines(message.content, 4)
             and not TokenRemover.find_token_in_message(message)
         )
 
diff --git a/bot/cogs/codeblock/parsing.py b/bot/cogs/codeblock/parsing.py
index 1bdb3b492..332a1deb0 100644
--- a/bot/cogs/codeblock/parsing.py
+++ b/bot/cogs/codeblock/parsing.py
@@ -5,6 +5,8 @@ import logging
 import re
 from typing import NamedTuple, Optional, Sequence
 
+from bot.utils import has_lines
+
 log = logging.getLogger(__name__)
 
 BACKTICK = "`"
@@ -82,7 +84,7 @@ def find_code_blocks(message: str) -> Optional[Sequence[CodeBlock]]:
         if groups["tick"] == BACKTICK and language:
             log.trace("Message has a valid code block with a language; returning None.")
             return None
-        elif len(groups["code"].split("\n", 3)) > 3:
+        elif has_lines(groups["code"], 4):
             code_block = CodeBlock(groups["code"], language, groups["tick"])
             code_blocks.append(code_block)
         else:
diff --git a/bot/utils/__init__.py b/bot/utils/__init__.py
index 5a6e1811b..4a02dc802 100644
--- a/bot/utils/__init__.py
+++ b/bot/utils/__init__.py
@@ -13,6 +13,14 @@ class CogABCMeta(CogMeta, ABCMeta):
     pass
 
 
+def has_lines(string: str, count: int) -> bool:
+    """Return True if `string` has at least `count` lines."""
+    split = string.split("\n", count - 1)
+
+    # Make sure the last part isn't empty, which would happen if there was a final newline.
+    return split[-1] and len(split) == count
+
+
 def pad_base64(data: str) -> str:
     """Return base64 `data` with padding characters to ensure its length is a multiple of 4."""
     return data + "=" * (-len(data) % 4)