Remove @ mentions

1 files changed, 7 insertions, 0 deletions

diff --git a/bot.py b/bot.py
index 495a22b..76f2952 100644
--- a/bot.py
+++ b/bot.py
@@ -54,10 +54,17 @@ if args.input is None:
         ])
     else:
         with open(args.data, 'r') as f:
+            # Get a line with at least two words
             lines = f.readlines()
             line = choice(lines).split()
             while len(line) < 2:
                 line = choice(lines).split()
+            
+            # Remove mentions
+            if line[0].count('@') > 1:
+                line[0] = '@'.join(line[0].split('@')[0:2])
+            if line[1].count('@') > 1:
+                line[1] = '@'.join(line[1].split('@')[0:2])
             args.input = line[0] + ' ' + line[1]