From c550dd2b610784ff36d2e20baf3601335daa192a Mon Sep 17 00:00:00 2001
From: Matthew Powers <matthewkevinpowers@gmail.com>
Date: Fri, 23 Feb 2024 10:51:08 -0500
Subject: [PATCH] format the keywords (#208)

---
 quinn/keyword_finder.py      | 22 +++++++++++++++++++++-
 tests/test_keyword_finder.py | 13 +++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/quinn/keyword_finder.py b/quinn/keyword_finder.py
index 4d89e88..b5cb83a 100644
--- a/quinn/keyword_finder.py
+++ b/quinn/keyword_finder.py
@@ -46,7 +46,7 @@ def search_file(path, keywords=default_keywords):
         for line_number, line in enumerate(f, 1):
             for keyword in keywords:    
                 if keyword in line:
-                    print(f"{line_number}: {line}", end='')
+                    print(f"{line_number}: {keyword_format(line)}", end='')
                     break
 
 
@@ -55,3 +55,23 @@ def search_files(path, keywords=default_keywords):
     file_list = [f for f in iglob(rootdir_glob, recursive=True) if os.path.isfile(f)]
     for f in file_list:
         search_file(f)
+
+
+def keyword_format(input, keywords=default_keywords):
+    nc = '\033[0m'
+    red = '\033[31m'
+    bold = '\033[1m'
+    res = input
+    for keyword in keywords:
+        res = surround_substring(res, keyword, red+bold, nc)
+    return res
+
+
+def surround_substring(input, substring, surround_start, surround_end):
+    index = input.find(substring)
+    res = ""
+    if index == -1:
+        res = input
+    else:
+        res = input[:index] + surround_start + substring + surround_end + input[(index+len(substring)):]
+    return res
diff --git a/tests/test_keyword_finder.py b/tests/test_keyword_finder.py
index a64db20..b02c6fd 100644
--- a/tests/test_keyword_finder.py
+++ b/tests/test_keyword_finder.py
@@ -1,8 +1,17 @@
 import pytest
-from quinn.keyword_finder import search_file, search_files
+from quinn.keyword_finder import search_file, search_files, keyword_format, surround_substring
 
 def test_search_file():
     search_file("tests/test_files/some_pyspark.py")
 
 def test_search_files():
-    search_files("tests/test_files")
\ No newline at end of file
+    search_files("tests/test_files")
+
+def test_keyword_format():
+    print(keyword_format("spark rdd stuff"))
+    print(keyword_format("spark rdd stuff with bad _jvm"))
+    print(keyword_format("nice string"))
+    print(keyword_format(""))
+
+def test_surround_substring():
+    print(surround_substring("spark rdd stuff", "rdd", "**", "||"))
\ No newline at end of file