From 9b2077780fafe173643bda91b7c97d0151336b98 Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Sun, 11 Dec 2022 23:50:44 -0500
Subject: [PATCH] Add interactive mode for picking sample inputs

---
 download.py | 58 ++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 7 deletions(-)

diff --git a/download.py b/download.py
index 14406ed..8e8f15a 100644
--- a/download.py
+++ b/download.py
@@ -7,7 +7,7 @@ import pathlib
 import sys
 import time
 import urllib.parse
-from typing import List, Optional
+from typing import Iterable, List, Optional
 import random
 
 import bs4
@@ -22,6 +22,7 @@ LOG = logging.getLogger(__name__)
 
 requests.utils.default_user_agent = lambda: "aoc-get by nick@ollien.com"
 
+
 @dataclasses.dataclass
 class PuzzleDate:
     year: int
@@ -85,7 +86,6 @@ def extract_sample_inputs(puzzle_source: str) -> List[str]:
         if pre_tag.code
     ]
 
-    LOG.info(f"Found {len(sample_inputs)} sample input(s)")
     return sample_inputs
 
 
@@ -96,10 +96,33 @@ def ensure_ends_with(s: str, c: str) -> str:
         return s
 
 
-def download_sample_inputs(date: PuzzleDate, output_dir: pathlib.Path):
+def find_sample_inputs(date: PuzzleDate) -> List[str]:
     LOG.debug("Downloading sample inputs")
     puzzle_source = get_puzzle_source(date)
     sample_inputs = extract_sample_inputs(puzzle_source)
+    LOG.info(f"Found {len(sample_inputs)} sample input(s)")
+
+    return sample_inputs
+
+
+def interactively_filter_sample_inputs(sample_inputs: List[str]) -> Iterable[str]:
+    def prompt_for_choice():
+        while True:
+            choice = input("Download as sample input? [Y/n] ")
+            if choice.lower() == "y" or not choice:
+                return True
+            elif choice.lower() == "n":
+                return False
+
+    for sample_input in sample_inputs:
+        print(sample_input, end="\n\n")
+        should_use = prompt_for_choice()
+        if should_use:
+            yield sample_input
+
+
+def save_sample_inputs(sample_inputs: Iterable[str], output_dir: pathlib.Path):
+    num_inputs = 0
     for i, sample_input in enumerate(sample_inputs, start=1):
         filename = f"sample-{i}.txt"
         path = output_dir / filename
@@ -108,8 +131,17 @@ def download_sample_inputs(date: PuzzleDate, output_dir: pathlib.Path):
             sample_file.write(data_to_write)
 
         LOG.debug(f"Wrote sample input {i} to {path}")
+        num_inputs += 1
 
-    LOG.info(f"Downloaded {len(sample_inputs)} sample input(s) successfully")
+    LOG.info(f"Downloaded {num_inputs} sample input(s) successfully")
+
+
+def download_sample_inputs(date: PuzzleDate, output_dir: pathlib.Path, *, interactive):
+    inputs = find_sample_inputs(date)
+    if interactive:
+        inputs = interactively_filter_sample_inputs(inputs)
+
+    save_sample_inputs(inputs, output_dir)
 
 
 def fetch_input(date: PuzzleDate, token: str) -> str:
@@ -198,7 +230,18 @@ class FetchCommand(click.Command):
     help="The location of the advent of code session token. Ignored if --token is provided.",
 )
 @click.option("--wait", is_flag=True, help="Wait until new puzzle inputs are ready")
-@click.option("--inputonly", "input_only", is_flag=True, help="Don't attempt to fetch sample inputs")
+@click.option(
+    "--inputonly",
+    "input_only",
+    is_flag=True,
+    help="Don't attempt to fetch sample inputs",
+)
+@click.option(
+    "-i",
+    "--interactive",
+    is_flag=True,
+    help="Interactively select sample inputs to download",
+)
 @click.option("--year", type=int)
 @click.option("--day", type=int)
 def main(
@@ -208,6 +251,7 @@ def main(
     output_dir: pathlib.Path,
     wait: bool,
     input_only: bool,
+    interactive: bool,
     day: Optional[int],
     year: Optional[int],
 ):
@@ -232,9 +276,9 @@ def main(
     date_to_fetch = specified_date or PuzzleDate.get_latest()
     step = None
     try:
+        step = "download sample inputs"
         if not input_only:
-            step = "download sample inputs"
-            download_sample_inputs(date_to_fetch, output_dir)
+            download_sample_inputs(date_to_fetch, output_dir, interactive=interactive)
         step = "download input"
         download_input(date_to_fetch, token, output_dir)
     except requests.exceptions.HTTPError as err: