note fix #8 and iterator over while trues (#10)

* note fix #8 and iterator over while trues

* color fixes

stacks_analyzer/detectors/AssertBlockHeight.py

@@ -15,11 +15,7 @@ def visit_node(self, node: Node, i):
         if i > 1:
             return
         if str(node.text, "utf8") == "asserts!":
-            descendants = NodeIterator(node.parent)
-            while True:
-                n = descendants.next()
-                if n is None:
-                    break
+            for n in NodeIterator(node.parent):
                 if str(n.text, "utf8") == "block-height" and n.grammar_name == "global":
                     pretty_print_warn(
                         self,

stacks_analyzer/detectors/CallInsideAsContract.py

@@ -17,11 +17,7 @@ def visit_node(self, node: Node, i):
         if i > 1:
             pass
         if str(node.text, "utf8") == "as-contract":
-            descendants = NodeIterator(node.parent)
-            while True:
-                n = descendants.next()
-                if n is None:
-                    break
+            for n in NodeIterator(node.parent):
                 if str(n.text, "utf8") == "contract-call?":
                     self.call = True
                 if n.grammar_name == "contract_principal_lit":

stacks_analyzer/detectors/DivideBeforeMultiply.py

@@ -15,11 +15,7 @@ def visit_node(self, node: Node, i):
         if i > 1:
             return
         if node.grammar_name == "native_identifier" and str(node.text, "utf8") == "*":
-            descendants = NodeIterator(node.parent)
-            while True:
-                n = descendants.next()
-                if n is None:
-                    break
+            for n in NodeIterator(node.parent):
                 if str(n.text, "utf8") == "/" and n.grammar_name == "native_identifier":
                     pretty_print_warn(
                         self,

stacks_analyzer/detectors/TxSenderInAssert.py

@@ -14,11 +14,7 @@ def visit_node(self, node: Node, i):
         if i > 1:
             return
         if str(node.text, "utf8") == "asserts!":
-            descendants = NodeIterator(node.parent)
-            while True:
-                n = descendants.next()
-                if n is None:
-                    break
+            for n in NodeIterator(node.parent):
                 if str(n.text, "utf8") == "tx-sender" and n.grammar_name == "global":
                     pretty_print_warn(
                         self,

stacks_analyzer/print_message.py

@@ -5,20 +5,17 @@
 
 from .visitor import Visitor
 
+tty = sys.stdout.isatty()
+
 
 class TerminalColors:
-    HEADER = '\033[95m'
-    OKBLUE = '\033[94m'
-    OKCYAN = '\033[96m'
-    OKGREEN = '\033[92m'
-    WARNING = '\033[93m'
-    FAIL = '\033[91m'
-    ENDC = '\033[0m'
-    BOLD = '\033[1m'
-    UNDERLINE = '\033[4m'
-
-
-def pretty_print_warn(visitor: Visitor, parent: Node, specific_node: Node, msg: str, help_msg: str | None, footnote: str | None):
+    HEADER = '\033[95m' if tty else ''
+    OKCYAN = '\033[96m' if tty else ''
+    ENDC = '\033[0m' if tty else ''
+
+
+def pretty_print_warn(visitor: Visitor, parent: Node, specific_node: Node, msg: str, help_msg: str | None,
+                      footnote: str | None):
     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]
@@ -28,28 +25,14 @@ def pretty_print_warn(visitor: Visitor, parent: Node, specific_node: Node, msg:
     arrows = "^" * (specific_node.end_point.column - specific_node.start_point.column)
     spaces = " " * ((specific_node.start_point.column * start_tabs) + 1)
 
-    tty = sys.stdout.isatty()
+    print(f"{TerminalColors.OKCYAN}Warning:{TerminalColors.ENDC} {msg}")
 
-    if tty:
-        print(f"{TerminalColors.OKCYAN}Warning:{TerminalColors.ENDC} {msg}")
-    else:
-        print(f"Warning: {msg}")
     print(f" {num_size_spaces}|")
     print(f" {line_number} | {contract_code}")
-    if tty:
-        print(f" {num_size_spaces}|{spaces}{TerminalColors.OKCYAN}{arrows}{TerminalColors.ENDC}")
-    else:
-        print(f" {num_size_spaces}|{spaces}{arrows}")
-
+    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:
-        if tty:
-            print(f" {num_size_spaces}{TerminalColors.OKCYAN}Note:{TerminalColors.ENDC}{footnote}")
-        else:
-            print(f" {num_size_spaces}Note: {footnote}")
-    print()
-
-
+        print(f" {num_size_spaces}{TerminalColors.OKCYAN}Note: {TerminalColors.ENDC}{footnote}")
 
+    print()

stacks_analyzer/stacks_analyzer.py

@@ -11,9 +11,7 @@
 from .detectors.UnwrapPanicUsage import UnwrapPanicUsage
 from .detectors.UpdatedFunctionsDetector import UpdatedFunctionsDetector
 from .detectors.VarCouldBeConstant import VarCouldBeConstant
-
 from .print_message import TerminalColors
-
 from .visitor import LinterRunner, Visitor
 
 DETECTOR_MAP = {
@@ -39,7 +37,6 @@ def main():
     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")
 
-
     user_args = arg_parser.parse_args()
     if user_args.command == "lint":
         filters = user_args.filter or list(DETECTOR_MAP.keys())
@@ -60,19 +57,12 @@ def main():
         max_length = max(len(st) for st in detectors)
         s = max_length // 2 - 4
 
-        if sys.stdout.isatty():
-            color = TerminalColors.OKCYAN
-            end = TerminalColors.ENDC
-        else:
-            color = ""
-            end = ""
+        print(f" {TerminalColors.OKCYAN}┌" + "─" * (s - 1) + " Detectors " + "─" * s + f"┐{TerminalColors.ENDC}")
+        for file in detectors:
+            print(
+                f" {TerminalColors.OKCYAN}|{TerminalColors.ENDC} {file.ljust(max_length + 1)}{TerminalColors.OKCYAN}|{TerminalColors.ENDC}")
+        print(f" {TerminalColors.OKCYAN}└" + "─" * (max_length + 2) + f"┘{TerminalColors.ENDC}")
 
-        if sys.stdout.isatty():
-            print(f" {color}┌" + "─" * (s - 1) + " Detectors " + "─" * s + f"┐{end}")
-            for file in detectors:
-                print(
-                    f" {color}|{end} {file.ljust(max_length + 1)}{color}|{end}")
-            print(f" {color}└" + "─" * (max_length + 2) + f"┘{end}")
 
 def get_detectors(filters: str, excludes):
     all_detectors = list(DETECTOR_MAP.keys())
@@ -99,14 +89,8 @@ def get_detectors(filters: str, excludes):
     return detectors
 
 
-
 def lint_file(path, lints: [Visitor]):
-    tty = sys.stdout.isatty()
-    if tty:
-        print(f"{TerminalColors.HEADER}====== Linting {path}... ======{TerminalColors.ENDC}")
-
-    else:
-        print(f"====== Linting {path}... ======")
+    print(f"{TerminalColors.HEADER}====== Linting {path}... ======{TerminalColors.ENDC}")
 
     with open(path, 'r') as file:
         source = file.read()

stacks_analyzer/visitor.py

@@ -35,7 +35,7 @@ def next(self) -> Node | None:
         while True:
             node = self.node()
 
-            if not self.visited.__contains__(node):
+            if node not in self.visited:
                 if self.cursor.goto_first_child():
                     continue
                 self.visited.append(node)
@@ -55,6 +55,15 @@ def next(self) -> Node | None:
     def node(self) -> Node | None:
         return self.cursor.node
 
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> Node | None:
+        node = self.next()
+        if node is None:
+            raise StopIteration
+        return node
+
 
 class LinterRunner:
     source: str