Add -c and -o options.

1 files changed, 50 insertions, 23 deletions

diff --git a/epubgrep.py b/epubgrep.py
index 78728d7..9118787 100755
--- a/epubgrep.py
+++ b/epubgrep.py
@@ -15,18 +15,19 @@ log = logging.getLogger('epubgrep')
 
 
 def get_chapter_title(mdata: List[Dict[str, Any]], fname: str) -> Optional[Tuple[str, int]]:
-    found_list = [(x['title'], x['index']) for x in mdata if x['src'] == fname]
-    log.debug('found_list = %s', found_list)
-    if len(found_list) > 0:
-        chap_title = found_list[0][0].strip(' \t.0123456789')
-        return chap_title, found_list[0][1]
-    else:
-        return ('Unknown', 0)
+    if mdata is not None:
+        found_list = [(x['title'], x['index']) for x in mdata if x['src'] == fname]
+        if len(found_list) > 0:
+            chap_title = found_list[0][0].strip(' \t.0123456789')
+            return chap_title, found_list[0][1]
+        else:
+            return ('Unknown', 0)
 
 
-def grep_book(filename: str, pattern: str, flags: int):
+def grep_book(filename: str, pattern: str, flags: int, counting: bool=False, color: bool=False):
     assert os.path.isfile(filename), "{} is not EPub file.".format(filename)
-    sought_RE = re.compile(pattern, flags)
+    sought_RE = re.compile('(' + pattern + ')', flags)
+    count = 0
 
     mline = flags & re.M == re.M
 
@@ -43,37 +44,62 @@ def grep_book(filename: str, pattern: str, flags: int):
                     if not printed_booktitle:
                         print('{}'.format(filename))
                         printed_booktitle = True
-                    chap_info = get_chapter_title(metadata.toc, zif.filename)
-                    print("{}. {}:\n".format(chap_info[1], chap_info[0]))
-                    print('{}\n'.format(res.group(0)))
+                    if counting:
+                        count += 1
+                    else:
+                        chap_info = get_chapter_title(metadata.toc, zif.filename)
+                        print("{}. {}:\n".format(chap_info[1], chap_info[0]))
+                        print('{}\n'.format(res.group(0)))
             else:
                 printed_title = False
                 for line in inf:
                     decoded_line = line.decode(errors='replace').strip()
-                    if sought_RE.search(decoded_line):
+                    res = sought_RE.search(decoded_line)
+                    if res:
                         if not printed_booktitle:
                             print('{}'.format(filename))
                             printed_booktitle = True
-                        if not printed_title:
-                            chap_info = get_chapter_title(metadata.toc,
-                                                          zif.filename)
-                            print("{}. {}:\n".format(chap_info[1], chap_info[0]))
-                            printed_title = True
-                        print(decoded_line)
+                        if counting:
+                            count += 1
+                        else:
+                            if not printed_title:
+                                chap_info = get_chapter_title(metadata.toc,
+                                                              zif.filename)
+                                if chap_info is not None:
+                                    print("{}. {}:\n".format(chap_info[1], chap_info[0]))
+                                    printed_title = True
+                        # https://stackoverflow.com/a/33206814
+                        # print("\033[31;1;4mHello\033[0m")
+                        if color:
+                            found_line = decoded_line.replace(
+                                res.group(1),
+                                "\033[31;1m" + res.group(1) + "\033[31;0m")
+                            print('{}\n'.format(found_line))
+                        else:
+                            print(decoded_line)
+
+    if count > 0:
+        print('Found: {}'.format(count))
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Grep through EPub book')
     parser.add_argument('pattern')
-    parser.add_argument('filename')
+    parser.add_argument('files', nargs='+')
+    parser.add_argument('-c', '--count',
+                        action='store_true',
+                        help="just counts of found patterns")
     parser.add_argument('-i', '--ignore-case',
                         action='store_true',
                         help="make search case insensitive")
+    parser.add_argument('-o', '--color', '--colour',
+                        action='store_true',
+                        help="make search case insensitive")
     parser.add_argument('-m', '--multi-line',
                         action='store_true',
                         help="make search multi line")
     args = parser.parse_args()
-    log.debug('args = %s', args)
+    # log.debug('args = %s', args)
 
     search_flags = 0
     if args.ignore_case:
@@ -82,5 +108,6 @@ if __name__ == "__main__":
     if args.multi_line:
         search_flags |= re.M | re.S
 
-    book_fname = os.path.realpath(args.filename)
-    grep_book(book_fname, args.pattern, search_flags)
+    for filename in args.files:
+        book_fname = os.path.realpath(filename)
+        grep_book(book_fname, args.pattern, search_flags, args.count, args.color)