Added -A, -B and -C flags for context in messages (#22)

Added -A, -B and -C flags for context in messages, with the same functionality as `grep` of giving more context on each find.

Author: faculerena

src/stacy_analyzer/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.0"
+__version__ = "1.0.0"
 
 from .analyzer import *
 

src/stacy_analyzer/analyzer.py

@@ -52,20 +52,39 @@ def main(self):
         lint_parser.add_argument("--exclude", nargs="+", type=str,
                                  help="Comma-separated list of detector names to exclude")
         list_detectors = subparsers.add_parser("detectors", help="List detectors")
-
+        lint_parser.add_argument("-A", nargs="+", type=int,
+                                 help="Print A lines of trailing context after findings.")
+        lint_parser.add_argument("-B", nargs="+", type=int,
+                                 help="Print B lines of leading context before findings.")
+        lint_parser.add_argument("-C", nargs="+", type=int,
+                                 help="Print C lines of leading and trailing context after and before findings. Takes precedence over -A and -B")
         user_args = arg_parser.parse_args()
         if user_args.command == "lint":
+
+            tc = 0
+            lc = 0
+
+            if user_args.A is not None:
+                tc = user_args.A[0]
+
+            if user_args.B is not None:
+                lc = user_args.B[0]
+
+            if user_args.C is not None:
+                tc = user_args.C[0]
+                lc = user_args.C[0]
+
             filters = list(self.DETECTOR_MAP.keys()) if user_args.filter is None else user_args.filter[0].split(',')
             excludes = [] if user_args.exclude is None else user_args.exclude[0].split(',')
             detectors = self.get_detectors(filters, excludes)
             path = user_args.path
             if path.endswith(".clar"):
-                self.lint_file(path, detectors, True)
+                self.lint_file(path, detectors, True, lc, tc)
             else:
                 for root, _, files in os.walk(path):
                     for file in files:
                         if file.endswith(".clar"):
-                            self.lint_file(os.path.join(root, file), detectors, True)
+                            self.lint_file(os.path.join(root, file), detectors, True, lc, tc)
 
         if user_args.command == "detectors":
             convert_camel_case = lambda s: s[0] + ''.join(' ' + c if c.isupper() else c for c in s[1:])
@@ -103,7 +122,8 @@ def get_detectors(self, filters: str, excludes: str):
 
         return [self.DETECTOR_MAP[name] for name in filtered_names]
 
-    def lint_file(self, filename, lints: [Visitor], print_output: bool):
+    def lint_file(self, filename, lints: [Visitor], print_output: bool, leading: int, trailing: int):
+
         if print_output:
             if self.isatty:
                 print(f"{TerminalColors.HEADER}====== Linting {filename}... ======{TerminalColors.ENDC}")
@@ -113,7 +133,7 @@ def lint_file(self, filename, lints: [Visitor], print_output: bool):
             source = file.read()
 
         runner: LinterRunner = LinterRunner(source, print_output, filename)
-        runner.add_lints(lints)
+        runner.add_lints(lints, leading, trailing)
 
         findings: [Finding] = runner.run()
 

src/stacy_analyzer/linter_runner.py

@@ -40,9 +40,10 @@ def add_lint(self, lint):
         self.lints.append(lint)
         return self
 
-    def add_lints(self, lint_classes: [Visitor]):
+    def add_lints(self, lint_classes: [Visitor], leading: int, trailing: int):
         for lint_class in lint_classes:
             lint = lint_class(self.print_output)
+            lint.set_context(leading, trailing)
             lint.add_source(self.source, self.src_name)
             self.lints.append(lint)
 

src/stacy_analyzer/print_message.py

@@ -11,28 +11,42 @@ class TerminalColors:
     OKCYAN = '\033[96m' if tty else ''
     ERROR = '\033[31;1;4m' if tty else ''
     BOLD = '\033[1m' if tty else ''
+    GREY = '\033[38;5;248m' if tty else ''
     ENDC = '\033[0m' if tty else ''
+    EXCYAN = ''
+
 
 
 def pretty_print_warn(visitor, parent: Node, specific_node: Node, msg: str, help_msg: str | None,
-                      footnote: str | None):
+                      footnote: str | None, leading_context: int, trailing_context: int):
     line_number = parent.start_point.row + 1
     num_size_spaces = " " * (int(math.log10(line_number)) + 2)
-    contract_code = visitor.source.split('\n')[line_number - 1]
-    start_tabs = contract_code.count('\t') + 1
-    contract_code = contract_code.replace('\t', '    ')
 
-    arrows = "^" * (specific_node.end_point.column - specific_node.start_point.column)
-    spaces = " " * ((specific_node.start_point.column * start_tabs) + 1)
+    all_lines = visitor.source.split('\n')
+    total_lines = len(all_lines)
+
+    start_line = max(0, line_number - leading_context - 1)
+    end_line = min(total_lines, line_number + trailing_context)
 
     print(f"{TerminalColors.OKCYAN}Warning:{TerminalColors.ENDC} {msg}")
+    print(f" {num_size_spaces}{TerminalColors.OKCYAN}|{TerminalColors.ENDC}")
+
+    for i in range(start_line, end_line):
+        current_line = i + 1
+        contract_code = all_lines[i]
+        start_tabs = contract_code.count('\t') + 1
+        contract_code = contract_code.replace('\t', '    ')
+
+        print(f" {TerminalColors.OKCYAN}{current_line} |{TerminalColors.ENDC} {contract_code}")
+
+        if current_line == line_number:
+            arrows = "^" * (specific_node.end_point.column - specific_node.start_point.column)
+            spaces = " " * ((specific_node.start_point.column * start_tabs) + 1)
+            print(f" {num_size_spaces}{TerminalColors.OKCYAN}|{TerminalColors.ENDC}{spaces}{TerminalColors.OKCYAN}{arrows}{TerminalColors.ENDC}")
+            if help_msg is not None:
+                print(f" {num_size_spaces}{TerminalColors.OKCYAN}|{TerminalColors.ENDC}{spaces}{help_msg}")
 
-    print(f" {num_size_spaces}|")
-    print(f" {line_number} | {contract_code}")
-    print(f" {num_size_spaces}|{spaces}{TerminalColors.OKCYAN}{arrows}{TerminalColors.ENDC}")
-    if help_msg is not None:
-        print(f" {num_size_spaces}|{spaces}{help_msg}")
     if footnote is not None:
-        print(f" {num_size_spaces}{TerminalColors.OKCYAN}Note: {TerminalColors.ENDC}{footnote}")
+        print(f" {num_size_spaces}{TerminalColors.OKCYAN}Note: {TerminalColors.GREY}{footnote}{TerminalColors.ENDC}")
 
     print()

src/stacy_analyzer/visitor.py

@@ -4,7 +4,8 @@
 from tree_sitter import Node
 
 from stacy_analyzer.print_message import pretty_print_warn
-
+LEADING_CONTEXT = 0
+TRAILING_CONTEXT = 0
 
 @dataclass
 class Location:
@@ -32,13 +33,19 @@ class Visitor:
     FOOTNOTE: str | None
     print_output: bool
     ignores: dict[str, ([int], Node)]
+    leading_context: int
+    trailing_context: int
 
-    def __init__(self, print_output: bool):
+    def __init__(self, print_output: bool, ):
         self.ignores = {}
         self.source = self.src_name = None
         self.findings = []
         self.print_output = print_output
 
+    def set_context(self, leading: int, trailing: int):
+        self.leading_context = leading
+        self.trailing_context = trailing
+
     def add_source(self, src: str, src_name: str = None):
         self.source = src
         self.src_name = src_name
@@ -58,7 +65,7 @@ def add_finding(self, node: Node, specific_node: Node):
             return
 
         if self.print_output:
-            pretty_print_warn(self, node, specific_node, self.MSG, self.HELP, self.FOOTNOTE)
+            pretty_print_warn(self, node, specific_node, self.MSG, self.HELP, self.FOOTNOTE, self.leading_context, self.trailing_context)
 
         parent = node.parent
         line_number = parent.start_point.row + 1

tests/test_module1.py

@@ -23,7 +23,7 @@ def test_detector_file(self):
                 a = Analyzer()
                 for detector in a.DETECTOR_MAP.values():
                     a = Analyzer()
-                    ret = a.lint_file(filename, [detector], False)
+                    ret = a.lint_file(filename, [detector], False, 0, 0)
                     is_vulnerable_path = "vulnerable" in filename
                     is_correct_detector = detector.__name__.lower() in filename.replace('_', '')
 
@@ -45,7 +45,7 @@ def test_profile_time(self):
         start = time.time()
         for _ in range(1000):
             a = Analyzer()
-            a.lint_file(filename, [detectorKlass for detectorKlass in lints], False)
+            a.lint_file(filename, [detectorKlass for detectorKlass in lints], False, 0, 0)
         end = time.time()
         print(f'Running 1000 times in `unused_arguments` test case. Took: {end - start:f}s')