Solve day 4 (Python)

q4.py

@@ -0,0 +1,79 @@
+from typing import Iterable
+
+WORD = "XMAS"
+def words_from_grid(grid, startx, starty) -> Iterable[str]:
+    num_rows = len(grid)
+    num_cols = len(grid[0])
+
+    for (dx, dy) in (
+        (1, 0), (-1, 0),
+        (0, 1), (0, -1),
+        (1, 1), (-1, 1),
+        (1, -1), (-1, -1)
+    ):
+        x = startx
+        y = starty
+        word = ""
+        for i in range(len(WORD)):
+            if x >= num_cols or x < 0:
+                continue
+            if y >= num_rows or y < 0:
+                continue
+
+            word += grid[y][x]
+            x += dx
+            y += dy
+
+        yield word
+
+def cross_words_from_grid(grid, startx, starty) -> tuple[str, str]:
+    num_rows = len(grid)
+    num_cols = len(grid[0])
+
+    output = []
+    for (dx, dy) in (
+        (1, 1), (1, -1)
+    ):
+        x = startx - dx
+        y = starty - dy
+        word = ""
+        for _ in range(3):
+            if x >= num_cols or x < 0:
+                continue
+            if y >= num_rows or y < 0:
+                continue
+
+            word += grid[y][x]
+            x += dx
+            y += dy
+        output.append(word)
+    # o: tuple[str, str] = 
+    return tuple(output)
+
+def task1(grid: list[str]) -> int:
+    count = 0
+    for y in range(len(grid)):
+        for x in range(len(grid[0])):
+            for word in words_from_grid(grid, x, y):
+                if word == WORD or word == str(reversed(WORD)):
+                    count += 1
+    return count
+
+def task2(grid: list[str]) -> int:
+    count = 0
+    for y in range(len(grid)):
+        for x in range(len(grid[0])):
+            leading, anti = cross_words_from_grid(grid, x, y)
+            if (leading == "MAS" or leading == "SAM") and (anti == "MAS" or anti == "SAM"):
+                count += 1
+    return count
+
+# I can't believe I got both parts first try... this should've been a very error-prone episode!
+def main():
+    with open("i4.txt") as f:
+        text = f.read().splitlines()
+    print(task1(text))
+    print(task2(text))
+
+if __name__ == "__main__":
+    main()