python-examples 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. python_examples-0.2.0/Algorithms/arrays/subsequence_conflict.py +35 -0
  2. python_examples-0.2.0/Algorithms/matrices/utils.py +59 -0
  3. python_examples-0.2.0/Algorithms/numerical_methods/coin_run.py +77 -0
  4. python_examples-0.2.0/Algorithms/numerical_methods/euler/p0.py +9 -0
  5. python_examples-0.2.0/Algorithms/numerical_methods/euler/p1.py +8 -0
  6. python_examples-0.2.0/Algorithms/numerical_methods/euler/p2.py +16 -0
  7. python_examples-0.2.0/Algorithms/numerical_methods/euler/p48.py +8 -0
  8. python_examples-0.2.0/Algorithms/numerical_methods/euler/p49.py +100 -0
  9. python_examples-0.2.0/Algorithms/numerical_methods/euler/p50.py +22 -0
  10. python_examples-0.2.0/Algorithms/numerical_methods/euler/p51.py +45 -0
  11. python_examples-0.2.0/Algorithms/numerical_methods/euler/p54.py +161 -0
  12. python_examples-0.2.0/Algorithms/numerical_methods/euler/utils.py +16 -0
  13. python_examples-0.2.0/Algorithms/slidingWindow/longestUnique.py +37 -0
  14. python_examples-0.2.0/Algorithms/static_analysis/ast_tree.py +76 -0
  15. python_examples-0.2.0/Algorithms/static_analysis/deduplicate_security_findings.py +22 -0
  16. python_examples-0.2.0/Algorithms/static_analysis/package_dependency.py +70 -0
  17. python_examples-0.2.0/Algorithms/static_analysis/rule_matcher.py +70 -0
  18. python_examples-0.2.0/Algorithms/tensors/broadcast.py +98 -0
  19. python_examples-0.2.0/Algorithms/tensors/kmeans.py +49 -0
  20. python_examples-0.2.0/LICENSE +674 -0
  21. python_examples-0.2.0/LLMs/torch_examples/batch_agg.py +78 -0
  22. python_examples-0.2.0/LLMs/torch_examples/pooling.py +26 -0
  23. python_examples-0.2.0/LLMs/torch_examples/shape_literacy.py +37 -0
  24. python_examples-0.2.0/LLMs/xai.py +92 -0
  25. python_examples-0.2.0/PKG-INFO +32 -0
  26. python_examples-0.2.0/README.md +22 -0
  27. python_examples-0.2.0/mixins/__init__.py +0 -0
  28. python_examples-0.2.0/mixins/benchmark.py +16 -0
  29. python_examples-0.2.0/mixins/logging.py +7 -0
  30. python_examples-0.2.0/pyproject.toml +29 -0
  31. python_examples-0.2.0/python_examples.egg-info/PKG-INFO +32 -0
  32. python_examples-0.2.0/python_examples.egg-info/SOURCES.txt +60 -0
  33. python_examples-0.2.0/python_examples.egg-info/dependency_links.txt +1 -0
  34. python_examples-0.2.0/python_examples.egg-info/requires.txt +1 -0
  35. python_examples-0.2.0/python_examples.egg-info/top_level.txt +3 -0
  36. python_examples-0.2.0/setup.cfg +4 -0
@@ -0,0 +1,35 @@
1
+ from typing import List
2
+
3
+ def dfs(start, adj_lst, visited, seen):
4
+ """
5
+ return: if cycle is detected
6
+ """
7
+ if start in seen:
8
+ return True
9
+ if start in visited or start not in adj_lst:
10
+ return False
11
+ seen.add(start)
12
+ for neighbor in adj_lst[start]:
13
+ if dfs(neighbor, adj_lst, visited, seen):
14
+ return True
15
+ seen.remove(start)
16
+ visited.add(start)
17
+ return False
18
+
19
+ def subsequence(input_lsts: List[List[int]]):
20
+ """
21
+ input_lsts list of lists
22
+ """
23
+ adj_lst = dict()
24
+
25
+ for l in input_lsts:
26
+ for i in range(1, len(l)):
27
+ adj_lst.setdefault(l[i-1], []).append(l[i])
28
+
29
+ visited = set()
30
+ for node in adj_lst.keys():
31
+ if node not in visited:
32
+ is_cycle = dfs(node, adj_lst, visited, set())
33
+ if is_cycle:
34
+ return False
35
+ return True
@@ -0,0 +1,59 @@
1
+ import numpy as np
2
+ from scipy.linalg.blas import dsymm
3
+ import time
4
+
5
+
6
+ def slow_update_matrix(C: np.ndarray, A: np.ndarray, B: np.ndarray):
7
+ n = C.shape[0]
8
+ for i in range(n):
9
+ for j in range(n):
10
+ for k in range(n):
11
+ C[i, j] += A[k, i] * B[k, j] * A[k, j]
12
+ return C
13
+
14
+
15
+ def update_matrix(C: np.ndarray, A: np.ndarray, B: np.ndarray):
16
+ """
17
+ Update matrix C by adding A^T @ B @ A, for symmetric matrix B efficiently
18
+ and ensure innermost loop uses unit-stride vector operations. Since B is symmetric,
19
+ it's stored in compact upper triangular form.
20
+ """
21
+ assert C.shape == A.shape and C.shape == B.shape, (
22
+ "Matrices C, A, and B must have the same shape."
23
+ )
24
+ assert (B.T == B).all(), "Matrix B must be symmetric."
25
+ BA = B @ A
26
+ C += A.T @ BA
27
+ return C
28
+
29
+
30
+ def symm_update(C, A, B):
31
+ BA = dsymm(1.0, B, A) # B symmetric
32
+ C += A.T @ BA
33
+ return C
34
+
35
+
36
+ if __name__ == "__main__":
37
+ n = 10000 # try larger like 400–1000 for clearer results
38
+ A = np.random.randn(n, n)
39
+ B = np.random.randn(n, n)
40
+ B = (B + B.T) / 2 # symmetrize
41
+ C1 = np.zeros((n, n))
42
+ C2 = np.zeros((n, n))
43
+
44
+ # Time slow
45
+ # start = time.time()
46
+ # slow_update_matrix(C1.copy(), A, B)
47
+ # print("Slow update time:", time.time() - start)
48
+
49
+ # Time fast
50
+ start = time.time()
51
+ update_matrix(C2.copy(), A, B)
52
+ print("Fast update time:", time.time() - start)
53
+
54
+ start = time.time()
55
+ symm_update(C2.copy(), A, B)
56
+ print("BLAS symm update time:", time.time() - start)
57
+
58
+ # Check correctness
59
+ print("Difference norm:", np.linalg.norm(C1 - C2))
@@ -0,0 +1,77 @@
1
+ import numpy as np
2
+
3
+ from mixins.benchmark import BenchmarkMixin
4
+ from mixins.logging import LoggingMixin
5
+
6
+
7
+ def count_max_length_runs(length: int, rng=None):
8
+ lst = np.random.randint(0, 2, length) if rng is None else rng.integers(0, 2, length)
9
+ i = 1
10
+ num_runs = 0
11
+ curr_max = 0
12
+ run = 1
13
+
14
+ while i < length:
15
+ if lst[i] == lst[i - 1]:
16
+ run += 1
17
+ else:
18
+ # only update number of runs if run ends
19
+ if run > curr_max:
20
+ curr_max = run
21
+ num_runs = 1
22
+ elif run == curr_max:
23
+ num_runs += 1
24
+ run = 1
25
+ i += 1
26
+ if run > curr_max:
27
+ curr_max = run
28
+ num_runs = 1
29
+ elif run == curr_max:
30
+ num_runs += 1
31
+ return num_runs
32
+
33
+
34
+ def count_num_runs(length: int, rng=None):
35
+ lst = np.random.randint(0, 2, length) if rng is None else rng.integers(0, 2, length)
36
+ runs = 1
37
+ for i in range(1, length):
38
+ if lst[i] != lst[i - 1]:
39
+ runs += 1
40
+ return runs
41
+
42
+
43
+ class CoinRunSimulation(BenchmarkMixin, LoggingMixin):
44
+ def __init__(self, seed=None):
45
+ self.rng = np.random.default_rng(seed)
46
+
47
+ def simulate(self, steps: int, length: int = 100):
48
+ if steps < 0:
49
+ raise ValueError("steps must be non-negative")
50
+ if length < 1:
51
+ raise ValueError("length must be positive")
52
+
53
+ results = []
54
+ for step in range(steps):
55
+ result = {
56
+ "step": step,
57
+ "max_length_runs": count_max_length_runs(length, self.rng),
58
+ "num_runs": count_num_runs(length, self.rng),
59
+ }
60
+ results.append(result)
61
+ self.logger.debug("Completed coin-run simulation step %s", step)
62
+ return results
63
+
64
+ def run(self, steps: int, length: int = 100):
65
+ return self.benchmark(self.simulate, steps, length)
66
+
67
+
68
+ def main():
69
+ simulation = CoinRunSimulation()
70
+ benchmark = simulation.run(steps=10_000, length=100)
71
+ results = benchmark["result"]
72
+ print(sum(result["max_length_runs"] for result in results) / len(results))
73
+ print(sum(result["num_runs"] for result in results) / len(results))
74
+
75
+
76
+ if __name__ == "__main__":
77
+ main()
@@ -0,0 +1,9 @@
1
+ def sum_squares(n: int):
2
+ if n < 1:
3
+ raise ValueError("n must be specified for first n sum squares")
4
+ s = 0
5
+ for i in range(1, n + 1):
6
+ v = i * i
7
+ if v % 2 != 0:
8
+ s += v
9
+ return s
@@ -0,0 +1,8 @@
1
+ def mult_35(n: int):
2
+ v = 1
3
+ s = 0
4
+ while v < n:
5
+ if v % 3 == 0 or v % 5 == 0:
6
+ s += v
7
+ v += 1
8
+ return s
@@ -0,0 +1,16 @@
1
+ def even_fib(n: int):
2
+ v1 = 1
3
+ v2 = 2
4
+ if n < 1:
5
+ raise ValueError("n must be at least 1")
6
+ elif n == 1:
7
+ return 0
8
+ else:
9
+ s = 0
10
+ while v2 <= n:
11
+ if v2 % 2 == 0:
12
+ s += v2
13
+ temp = v2
14
+ v2 = v2 + v1
15
+ v1 = temp
16
+ return s
@@ -0,0 +1,8 @@
1
+ def self_powers(n: int):
2
+ if n < 1:
3
+ raise ValueError("n must be at least 1")
4
+ res = 0
5
+ MOD = 10**10
6
+ for i in range(1, n + 1):
7
+ res = (res + pow(i, i, MOD)) % MOD
8
+ return res
@@ -0,0 +1,100 @@
1
+ from collections import defaultdict
2
+ import math
3
+
4
+
5
+ def is_prime(n: int):
6
+ if n < 1:
7
+ raise ValueError("n must be at least 1")
8
+ elif n == 1:
9
+ return False
10
+ elif n == 2:
11
+ return True
12
+ elif n % 2 == 0:
13
+ return False
14
+ else:
15
+ for i in range(3, int(math.sqrt(n)) + 1):
16
+ if n % i == 0:
17
+ return False
18
+ return True
19
+
20
+
21
+ def sieve_n_digit_primes(n: int):
22
+ if n < 1:
23
+ return []
24
+
25
+ low = 10 ** (n - 1)
26
+ high = 10**n - 1
27
+
28
+ if n == 1:
29
+ low = 2
30
+
31
+ limit = math.isqrt(high)
32
+ is_prime = [True] * (limit + 1)
33
+ base_primes = []
34
+
35
+ # sieve base primes
36
+ for p in range(2, limit + 1):
37
+ if is_prime[p]:
38
+ base_primes.append(p)
39
+ for i in range(p * p, limit + 1, p):
40
+ is_prime[i] = False
41
+
42
+ # Segment the range [low, high] into cache-friendly block sizes
43
+ # 32KB to 256KB block size prevents CPU cache thrashing
44
+ block_size = 500000
45
+ n_digit_primes = []
46
+
47
+ for current_low in range(low, high + 1, block_size):
48
+ current_high = min(current_low + block_size - 1, high)
49
+ range_size = current_high - current_low + 1
50
+
51
+ segment = [True] * range_size
52
+
53
+ for p in base_primes:
54
+ # Find the first multiple of p >= current_low and >= p^2
55
+ start_multiple = max(p * p, ((current_low + p - 1) // p) * p)
56
+
57
+ for j in range(start_multiple, current_high + 1, p):
58
+ segment[j - current_low] = False
59
+
60
+ for i in range(range_size):
61
+ if segment[i]:
62
+ n_digit_primes.append(current_low + i)
63
+
64
+ return n_digit_primes
65
+
66
+
67
+ def prime_permutations(n: int, perms: int) -> str:
68
+ """
69
+ n: number of digits in the prime permutations
70
+ end in odd
71
+ perms: number of permutations to consider
72
+ _ _ _ 1,3,5,7,9
73
+ """
74
+ all_primes = sieve_n_digit_primes(n)
75
+ table = defaultdict(list)
76
+ for prime in all_primes:
77
+ key = "".join(sorted(str(prime)))
78
+ table[key].append(prime)
79
+
80
+ res = []
81
+ for key, nums in table.items():
82
+ if len(nums) < perms:
83
+ continue
84
+ nums_set = set(nums)
85
+ for start in nums_set:
86
+ for next_num in nums_set:
87
+ if next_num <= start:
88
+ continue
89
+ diff = next_num - start
90
+ seq = [start]
91
+ curr = start
92
+ while curr in nums_set and len(seq) < perms:
93
+ curr += diff
94
+ if curr in nums_set:
95
+ seq.append(curr)
96
+ else:
97
+ break
98
+ if len(seq) == perms:
99
+ res.append("".join([str(v) for v in seq]))
100
+ return res
@@ -0,0 +1,22 @@
1
+ from typing import Optional
2
+ from Algorithms.numerical_methods.euler import utils
3
+
4
+
5
+ def consecutive_prime_sum(n: int) -> Optional[int]:
6
+ is_prime, primes = utils.sieve(n)
7
+ longest_sum = 0
8
+ longest_sum_prime = None
9
+ pref_sums = [0] * (len(primes) + 1)
10
+ # go through every prefix sum (pairs of start and end primes), sum should be prime
11
+ for i in range(1, len(pref_sums)):
12
+ pref_sums[i] = pref_sums[i - 1] + primes[i - 1]
13
+
14
+ for start in range(len(primes)):
15
+ for end in range(start + longest_sum + 1, len(primes) + 1):
16
+ s = pref_sums[end] - pref_sums[start]
17
+ if s >= n:
18
+ break
19
+ if is_prime[s]:
20
+ longest_sum = end - start
21
+ longest_sum_prime = s
22
+ return longest_sum_prime
@@ -0,0 +1,45 @@
1
+ from Algorithms.numerical_methods.euler import utils
2
+ from typing import Optional
3
+ from itertools import combinations
4
+
5
+
6
+ def prime_digit_replacements(
7
+ n: int, prime_family: int, start_point: Optional[int] = None
8
+ ) -> Optional[int]:
9
+ """
10
+ n: int
11
+ prime_family: int
12
+ return: Smallest prime that has a prime family of size `prime_family`
13
+ """
14
+ is_prime, primes = utils.sieve(n)
15
+ nums = "0123456789"
16
+ for p in primes:
17
+ if start_point is not None and p < start_point:
18
+ continue
19
+ s = str(p)
20
+ for digit in nums:
21
+ positions = [i for i, c in enumerate(s) if c == digit]
22
+
23
+ if not positions:
24
+ continue
25
+ for r in range(1, len(positions) + 1):
26
+ for subset in combinations(positions, r):
27
+ if subset[-1] == n - 1:
28
+ continue
29
+ count = 0
30
+ smallest = None
31
+ for replacement in nums:
32
+ if subset[0] == 0 and replacement == "0":
33
+ continue
34
+ updated = list(s)
35
+ for i in subset:
36
+ updated[i] = replacement
37
+
38
+ candidate = int("".join(updated))
39
+ if is_prime[candidate]:
40
+ count += 1
41
+ if smallest is None or candidate < smallest:
42
+ smallest = candidate
43
+ if count == prime_family:
44
+ return smallest
45
+ return None
@@ -0,0 +1,161 @@
1
+ from collections import defaultdict
2
+ from enum import IntEnum
3
+ from typing import Any, Dict, List
4
+
5
+ VALUE_MAP = {"T": 10, "J": 11, "Q": 12, "K": 13, "A": 14}
6
+
7
+
8
+ class Rank(IntEnum):
9
+ HIGH_CARD = 0
10
+ ONE_PAIR = 1
11
+ TWO_PAIRS = 2
12
+ THREE_OF_A_KIND = 3
13
+ STRAIGHT = 4
14
+ FLUSH = 5
15
+ FULL_HOUSE = 6
16
+ FOUR_OF_A_KIND = 7
17
+ STRAIGHT_FLUSH = 8
18
+ ROYAL_FLUSH = 9
19
+
20
+
21
+ def union_find_longest_consecutive(nums):
22
+ parent = {}
23
+ size = {}
24
+ for x in nums:
25
+ parent[x] = x
26
+ size[x] = 1
27
+
28
+ def find(x):
29
+ if parent[x] != x:
30
+ parent[x] = find(parent[x])
31
+ return parent[x]
32
+
33
+ def union(a, b):
34
+ ra, rb = find(a), find(b)
35
+ if ra == rb:
36
+ return
37
+ if size[ra] < size[rb]:
38
+ ra, rb = rb, ra
39
+ # assume rb is the smaller one
40
+ parent[rb] = ra
41
+ size[ra] += size[rb]
42
+
43
+ snums = set(nums)
44
+ for x in snums:
45
+ if x + 1 in snums:
46
+ union(x, x + 1)
47
+ return max(size[find(x)] for x in nums)
48
+
49
+
50
+ def gather_hand_information(hand: List[str]) -> Dict[str, Any]:
51
+ """
52
+ return:
53
+ - all_same_suit
54
+ - value counts
55
+ - are consecutive
56
+ """
57
+ seen = set()
58
+ all_same = True
59
+ d = defaultdict(int)
60
+ vals = []
61
+ for val, suit in [tuple(card) for card in hand]:
62
+ if suit not in seen:
63
+ if not seen:
64
+ seen.add(suit)
65
+ else:
66
+ all_same = False
67
+ int_val = VALUE_MAP[val] if val in VALUE_MAP else int(val)
68
+ d[int_val] += 1
69
+ vals.append(int_val)
70
+ all_consecutive = union_find_longest_consecutive(vals) == len(hand)
71
+ return {
72
+ "all_same_suit": all_same,
73
+ "value_counts": d,
74
+ "all_consecutive": all_consecutive,
75
+ }
76
+
77
+
78
+ def determine_rank(player_info: Dict[str, Any]) -> tuple[Rank, tuple]:
79
+ counts = player_info["value_counts"]
80
+
81
+ # values sorted by (count, value)
82
+ groups = sorted(((cnt, val) for val, cnt in counts.items()), reverse=True)
83
+
84
+ values_desc = sorted(counts.keys(), reverse=True)
85
+
86
+ is_flush = player_info["all_same_suit"]
87
+ is_straight = player_info["all_consecutive"]
88
+
89
+ # Handle A2345 straight
90
+ if set(counts.keys()) == {14, 2, 3, 4, 5}:
91
+ is_straight = True
92
+ straight_high = 5
93
+ elif is_straight:
94
+ straight_high = max(counts)
95
+ else:
96
+ straight_high = None
97
+
98
+ if is_flush and is_straight:
99
+ if straight_high == 14 and min(counts) == 10:
100
+ return Rank.ROYAL_FLUSH, ()
101
+ return Rank.STRAIGHT_FLUSH, (straight_high,)
102
+
103
+ if groups[0][0] == 4:
104
+ quad = groups[0][1]
105
+ kicker = groups[1][1]
106
+ return Rank.FOUR_OF_A_KIND, (quad, kicker)
107
+
108
+ if groups[0][0] == 3 and groups[1][0] == 2:
109
+ return Rank.FULL_HOUSE, (groups[0][1], groups[1][1])
110
+
111
+ if is_flush:
112
+ return Rank.FLUSH, tuple(values_desc)
113
+
114
+ if is_straight:
115
+ return Rank.STRAIGHT, (straight_high,)
116
+
117
+ if groups[0][0] == 3:
118
+ trip = groups[0][1]
119
+ kickers = sorted(
120
+ (v for v, c in counts.items() if c == 1),
121
+ reverse=True,
122
+ )
123
+ return Rank.THREE_OF_A_KIND, (trip, *kickers)
124
+
125
+ if groups[0][0] == 2 and groups[1][0] == 2:
126
+ pairs = sorted(
127
+ (v for v, c in counts.items() if c == 2),
128
+ reverse=True,
129
+ )
130
+ kicker = next(v for v, c in counts.items() if c == 1)
131
+ return Rank.TWO_PAIRS, (*pairs, kicker)
132
+
133
+ if groups[0][0] == 2:
134
+ pair = groups[0][1]
135
+ kickers = sorted(
136
+ (v for v, c in counts.items() if c == 1),
137
+ reverse=True,
138
+ )
139
+ return Rank.ONE_PAIR, (pair, *kickers)
140
+
141
+ return Rank.HIGH_CARD, tuple(values_desc)
142
+
143
+
144
+ def poker_hands(poker_file: str) -> int:
145
+ player1_wins = 0
146
+ with open(poker_file) as f:
147
+ for row in f.readlines():
148
+ vals = row.strip().split(" ")
149
+ player1, player2 = vals[:5], vals[5:]
150
+ player1_info = gather_hand_information(player1)
151
+ player2_info = gather_hand_information(player2)
152
+ rank1 = determine_rank(player1_info)
153
+ rank2 = determine_rank(player2_info)
154
+ if rank1 > rank2:
155
+ player1_wins += 1
156
+ elif rank1 == rank2:
157
+ if max(player1_info["value_counts"].keys()) > max(
158
+ player2_info["value_counts"].keys()
159
+ ):
160
+ player1_wins += 1
161
+ return player1_wins
@@ -0,0 +1,16 @@
1
+ # utils.py
2
+
3
+
4
+ def sieve(n: int):
5
+ is_prime = [True] * (n + 1)
6
+ if n >= 0:
7
+ is_prime[0] = False
8
+ if n >= 1:
9
+ is_prime[1] = False
10
+ primes = []
11
+ for p in range(2, n + 1):
12
+ if is_prime[p]:
13
+ primes.append(p)
14
+ for v in range(p * p, n + 1, p):
15
+ is_prime[v] = False
16
+ return is_prime, primes
@@ -0,0 +1,37 @@
1
+ from collections import OrderedDict
2
+ from pydantic import BaseModel, Field
3
+ from typing import List
4
+
5
+
6
+ class SlidingWindowParam(BaseModel):
7
+ s: str = Field(..., description="Input string", min_length=1)
8
+
9
+
10
+ def lengthOfLongestSubstringTwoDistinct(s: SlidingWindowParam) -> int:
11
+ seen = OrderedDict()
12
+ start = 0
13
+ longest = 0
14
+ for end in range(len(s)):
15
+ if s[end] in seen:
16
+ seen.move_to_end(s[end])
17
+ seen[s[end]] = end
18
+ if len(seen) > 2:
19
+ _, idx = seen.popitem(last=False)
20
+ start = idx + 1
21
+ longest = max(longest, end - start + 1)
22
+ return longest
23
+
24
+
25
+ def maximumUniqueSubarray(nums: List[int]) -> int:
26
+ last_seen = dict()
27
+ start, curr_sum, max_sum = 0, 0, 0
28
+ for end in range(len(nums)):
29
+ if nums[end] in last_seen:
30
+ last_idx = last_seen[nums[end]]
31
+ if last_idx >= start:
32
+ curr_sum -= sum(nums[start : (last_idx + 1)])
33
+ start = last_idx + 1
34
+ curr_sum += nums[end]
35
+ max_sum = max(max_sum, curr_sum)
36
+ last_seen[nums[end]] = end
37
+ return max_sum
@@ -0,0 +1,76 @@
1
+ from typing import Dict
2
+
3
+
4
+ class Node:
5
+ def __init__(self, node_type, value=None, children=None):
6
+ self.node_type = node_type
7
+ self.value = value
8
+ self.children = children if children is not None else []
9
+
10
+ def __repr__(self):
11
+ if self.value:
12
+ return f"{self.node_type}({self.value})"
13
+ return f"{self.node_type}({self.children})"
14
+
15
+ def __eq__(self, other):
16
+ return (
17
+ self.node_type == other.node_type
18
+ and self.value == other.value
19
+ and self.children == other.children
20
+ )
21
+
22
+
23
+ def find_calls(root: Node, function_name: str):
24
+ """
25
+ This function takes AST root and function name as input and returns all function calls in the AST.
26
+ First child node is the name of the function
27
+ """
28
+ calls = []
29
+
30
+ def gather_calls(root, function_name):
31
+ if root:
32
+ if root.node_type == "Call":
33
+ if (
34
+ root.children
35
+ and root.children[0].node_type == "Name"
36
+ and root.children[0].value == function_name
37
+ ):
38
+ calls.append(root)
39
+ for child in root.children:
40
+ gather_calls(child, function_name)
41
+
42
+ gather_calls(root, function_name)
43
+ return calls
44
+
45
+
46
+ def match(rule_node: Node, code_node: Node, captures: Dict[str, Node]):
47
+ """
48
+ Returns True if code_node matches rule_node.
49
+
50
+ captures should be populated with metavariable bindings.
51
+ """
52
+ if rule_node is None or code_node is None:
53
+ return rule_node is code_node
54
+
55
+ if rule_node.node_type == "MetaVar":
56
+ name = rule_node.value
57
+ if name in captures:
58
+ return captures[name] == code_node
59
+ captures[name] = code_node
60
+ return True
61
+
62
+ if rule_node.node_type != code_node.node_type:
63
+ return False
64
+
65
+ if rule_node.value is not None:
66
+ if code_node.value is None or rule_node.value != code_node.value:
67
+ return False
68
+
69
+ if len(rule_node.children) != len(code_node.children):
70
+ return False
71
+
72
+ for rule_child, code_child in zip(rule_node.children, code_node.children):
73
+ if not match(rule_child, code_child, captures):
74
+ return False
75
+
76
+ return True