Improve argument parsing and add delete line option

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = subprompt
-version = 0.0.4
+version = 0.0.5
 author = Augusto Lenz Gunsch
 author_email = augusto@augustogunsch.com
 description = Substitute Regex in files with prompt confirmation

src/subprompt/subprompt.py

@@ -25,52 +25,80 @@
 import re
 import hashlib
 import os
+import argparse
 from sys import argv, stderr
 from colorama import init, Fore, Back, Style
 
 init(autoreset=True)
 
 class Filter:
-    def __init__(self, regex, sub):
+    def __init__(self, args):
         self.replace_all = False
         self.quit_loop = False
         self.sub_count = 0
         self.match_count = 0
 
-        self.regex = re.compile(regex)
-        self.sub = sub
+        self.regex = re.compile(args.regex)
+        self.sub = args.r
+        self.delete = args.d
+        self.spread = args.n
 
-    def _prompt(self, matchobj, line, line_n, fname):
+        self.lines_to_delete = []
+
+    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:
             return curr_match
 
-        replaced_match = self.regex.sub(self.sub, curr_match)
+        if not self.delete:
+            replaced_match = self.regex.sub(self.sub, curr_match)
 
-        if replaced_match == curr_match:
-            return replaced_match
+            if replaced_match == curr_match:
+                return replaced_match
 
         self.match_count += 1
 
         if self.replace_all:
             self.sub_count += 1
+            if self.delete:
+                self.lines_to_delete.append(line_n)
+                return None
             return replaced_match
 
+        start_n = max(0, line_n - self.spread)
+        end_n = min(len(lines), line_n + self.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
+
+        if not self.delete:
+            cut_replaced[line_n] = line.replace(curr_match,
+                    Back.YELLOW + Fore.BLACK + replaced_match + Back.RESET + Fore.RESET)
+        else:
+            cut_replaced.pop(line_n)
+
+        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()
 
@@ -80,6 +108,9 @@ class Filter:
             if answer in ['y', 'yes', '']:
                 self.sub_count += 1
                 print()
+                if self.delete:
+                    self.lines_to_delete.append(line_n)
+                    return None
                 return replaced_match
 
             elif answer in ['n', 'no']:
@@ -89,6 +120,9 @@ class Filter:
             elif answer in ['a', 'all']:
                 self.sub_count += 1
                 self.replace_all = True
+                if self.delete:
+                    self.lines_to_delete.append(line_n)
+                    return None
                 return replaced_match
 
             elif answer in ['q', 'quit']:
@@ -105,12 +139,15 @@ 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)
                 ]
 
+        for line in reversed(self.lines_to_delete):
+            lines.pop(line)
+
         new_contents = '\n'.join(lines)
 
         hash_old = hashlib.md5(contents.encode())
@@ -122,23 +159,27 @@ class Filter:
                 file.write('\n')
 
 def run(args):
-    if len(args) < 3:
-        print('usage: subprompt [REGEX] [SUB] [FILES...]', file=stderr)
-        exit(1)
+    parser = argparse.ArgumentParser(description='Modifies lines matched by a regex interactively')
+    parser.add_argument('regex', metavar='REGEX', type=str)
+    action = parser.add_mutually_exclusive_group(required=True)
+    action.add_argument('-d', help='delete line', action='store_true')
+    action.add_argument('-r', help='replace match with expression', type=str)
+    parser.add_argument('files', metavar='FILES', nargs='+', type=str)
+    parser.add_argument('-n', help='size lines preview (default=3)', type=int, default=3)
 
-    regex = args[0]
-    sub = args[1]
-    files = args[2:]
+    args = parser.parse_args(args)
 
-    filter_obj = Filter(regex, sub)
+    filter_obj = Filter(args)
 
-    for file in files:
+    for file in args.files:
         # Check if file is writable
         if os.access(file, os.W_OK):
             filter_obj.filter_file(file)
 
             if filter_obj.quit_loop:
                 break
+        else:
+            print('WARNING: File {0} is not writable. Skipping...'.format(file))
 
     if filter_obj.match_count == 0:
         print(Fore.RED + Style.BRIGHT + 'No matches found.')
@@ -148,8 +189,7 @@ def run(args):
                         Fore.WHITE + str(filter_obj.match_count) + Fore.YELLOW))
 
 def main():
-    argv.pop(0)
-    run(argv)
+    run(argv[1:])
 
 if __name__ == '__main__':
     main()