From 8969ea895fc2ad14da8a32efb07a85d884693a55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Capelle?= <capelle.mikael@gmail.com>
Date: Sat, 14 Dec 2024 11:52:32 +0100
Subject: [PATCH] Handle PNG file generation.

---
 src/holt59/aoc/2024/day14.py      |  4 +++-
 src/holt59/aoc/base.py            | 16 +++++++++++++++-
 src/holt59/aoc/utils/api/files.py |  3 ++-
 src/holt59/aoc/utils/files.py     |  4 +++-
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/holt59/aoc/2024/day14.py b/src/holt59/aoc/2024/day14.py
index 2bded81..00e9fb6 100644
--- a/src/holt59/aoc/2024/day14.py
+++ b/src/holt59/aoc/2024/day14.py
@@ -3,6 +3,7 @@ import operator as op
 from math import prod
 from typing import Any, Iterator
 
+import numpy as np
 import parse  # pyright: ignore[reportMissingTypeStubs]
 
 from ..base import BaseSolver
@@ -64,8 +65,9 @@ class Solver(BaseSolver):
                             ).encode(),
                             True,
                         )
+                        self.files.image(f"result_{rnd+1}.png", np.array(m))
                     yield rnd + 1
-                    found = True
+                    # found = True
                     break
 
             if found:
diff --git a/src/holt59/aoc/base.py b/src/holt59/aoc/base.py
index 50cbaca..4d5b836 100644
--- a/src/holt59/aoc/base.py
+++ b/src/holt59/aoc/base.py
@@ -2,6 +2,8 @@ from abc import abstractmethod
 from logging import Logger
 from typing import Any, Final, Iterable, Iterator, Protocol, Sequence, TypeVar, overload
 
+from numpy.typing import NDArray
+
 _T = TypeVar("_T")
 
 
@@ -13,9 +15,21 @@ class ProgressHandler(Protocol):
     def wrap(self, values: Iterable[_T], total: int) -> Iterator[_T]: ...
 
 
-class FileHandler(Protocol):
+class FileHandler:
+    @abstractmethod
     def create(self, filename: str, content: bytes, text: bool = False): ...
 
+    def image(self, filename: str, image: NDArray[Any]):
+        from io import BytesIO
+
+        import matplotlib.pyplot as plt
+
+        io = BytesIO()
+        plt.imsave(io, image)  # type: ignore
+        io.seek(0)
+
+        self.create(filename, io.read(), False)
+
 
 class BaseSolver:
     def __init__(
diff --git a/src/holt59/aoc/utils/api/files.py b/src/holt59/aoc/utils/api/files.py
index d09d6ad..3d6650a 100644
--- a/src/holt59/aoc/utils/api/files.py
+++ b/src/holt59/aoc/utils/api/files.py
@@ -1,10 +1,11 @@
 from pathlib import Path
 from typing import Final
 
+from ...base import FileHandler
 from .base import dump_api_message
 
 
-class FileHandlerAPI:
+class FileHandlerAPI(FileHandler):
     def __init__(self, folder: Path):
         self.folder: Final = folder
 
diff --git a/src/holt59/aoc/utils/files.py b/src/holt59/aoc/utils/files.py
index 91c71ff..b7148ee 100644
--- a/src/holt59/aoc/utils/files.py
+++ b/src/holt59/aoc/utils/files.py
@@ -2,8 +2,10 @@ import logging
 from pathlib import Path
 from typing import Final
 
+from ..base import FileHandler
 
-class SimpleFileHandler:
+
+class SimpleFileHandler(FileHandler):
     def __init__(self, logger: logging.Logger, folder: Path):
         self.logger: Final = logger
         self.folder: Final = folder