Add surrouding lines preview

src/subprompt/subprompt.py

@@ -30,6 +30,9 @@ from colorama import init, Fore, Back, Style
 
 init(autoreset=True)
 
+# Config
+line_spread = 3
+
 class Filter:
     def __init__(self, regex, sub):
         self.replace_all = False
@@ -40,7 +43,14 @@ class Filter:
         self.regex = re.compile(regex)
         self.sub = sub
 
-    def _prompt(self, matchobj, line, line_n, fname):
+    def _number_lines(_, lines, start_n):
+        return [Fore.YELLOW + Style.BRIGHT +
+                str(n + start_n) +
+                Fore.RESET + Style.RESET_ALL +
+                ':' + line
+                for (n, line) in enumerate(lines)]
+
+    def _prompt(self, matchobj, lines, line, line_n, fname):
         curr_match = matchobj.group(0)
 
         if self.quit_loop:
@@ -57,20 +67,28 @@ class Filter:
             self.sub_count += 1
             return replaced_match
 
+        start_n = max(0, line_n - line_spread)
+        end_n = min(len(lines), line_n + line_spread)
+        cut = lines[start_n:end_n]
+
         highlighted = line.replace(curr_match,
                 Back.RED + curr_match + Back.RESET)
         replaced = line.replace(curr_match,
                 Back.YELLOW + Fore.BLACK + replaced_match + Back.RESET + Fore.RESET)
 
+        cut_highlighted = cut
+        cut_highlighted[line_n] = highlighted
+        cut_highlighted = self._number_lines(cut_highlighted, start_n)
+
+        cut_replaced = cut
+        cut_replaced[line_n] = replaced
+        cut_replaced = self._number_lines(cut_replaced, start_n)
+
         print(Fore.GREEN + Style.BRIGHT + fname)
 
-        print(Fore.YELLOW + Style.BRIGHT + str(line_n), end='')
-        print(':{0}'.format(highlighted))
-
+        print('\n'.join(cut_highlighted))
         print('Becomes the following:')
-
-        print(Fore.YELLOW + Style.BRIGHT + str(line_n), end='')
-        print(':{0}'.format(replaced))
+        print('\n'.join(cut_replaced))
 
         print()
 
@@ -105,7 +123,7 @@ class Filter:
 
         lines = [
                     self.regex.sub(
-                        lambda matchobj: self._prompt(matchobj, line, line_n, fname),
+                        lambda matchobj: self._prompt(matchobj, lines, line, line_n, fname),
                         line
                     )
                     for (line_n, line) in enumerate(lines)