Improve performance of linerange

bandit/core/utils.py

@@ -198,28 +198,64 @@ def escaped_bytes_representation(b):
     return b.decode("unicode_escape").encode("unicode_escape")
 
 
+def calc_linerange(node):
+    """Calculate linerange for subtree"""
+    if hasattr(node, "_bandit_linerange"):
+        return node._bandit_linerange
+
+    lines_min = 9999999999
+    lines_max = -1
+    if hasattr(node, "lineno"):
+        lines_min = node.lineno
+        lines_max = node.lineno
+    for n in ast.iter_child_nodes(node):
+        lines_minmax = calc_linerange(n)
+        lines_min = min(lines_min, lines_minmax[0])
+        lines_max = max(lines_max, lines_minmax[1])
+
+    node._bandit_linerange = (lines_min, lines_max)
+
+    return (lines_min, lines_max)
+
+
 def linerange(node):
     """Get line number range from a node."""
-    strip = {"body": None, "orelse": None, "handlers": None, "finalbody": None}
+    if hasattr(node, "_bandit_linerange_stripped"):
+        lines_minmax = node._bandit_linerange_stripped
+        return list(range(lines_minmax[0], lines_minmax[1] + 1))
+
+    strip = {
+        "body": None,
+        "orelse": None,
+        "handlers": None,
+        "finalbody": None,
+    }
     for key in strip.keys():
         if hasattr(node, key):
             strip[key] = getattr(node, key)
-            node.key = []
+            setattr(node, key, [])
 
     lines_min = 9999999999
     lines_max = -1
-    for n in ast.walk(node):
-        if hasattr(n, "lineno"):
-            lines_min = min(lines_min, n.lineno)
-            lines_max = max(lines_max, n.lineno)
+    if hasattr(node, "lineno"):
+        lines_min = node.lineno
+        lines_max = node.lineno
+    for n in ast.iter_child_nodes(node):
+        lines_minmax = calc_linerange(n)
+        lines_min = min(lines_min, lines_minmax[0])
+        lines_max = max(lines_max, lines_minmax[1])
 
     for key in strip.keys():
         if strip[key] is not None:
-            node.key = strip[key]
+            setattr(node, key, strip[key])
+
+    if lines_max == -1:
+        lines_min = 0
+        lines_max = 1
+
+    node._bandit_linerange_stripped = (lines_min, lines_max)
 
-    if lines_max > -1:
-        return list(range(lines_min, lines_max + 1))
-    return [0, 1]
+    return list(range(lines_min, lines_max + 1))
 
 
 def linerange_fix(node):