add block keyword to get_data

aocd/exceptions.py

@@ -9,5 +9,9 @@ class AocdError(Exception):
     """base exception for this package"""
 
 
+class PuzzleLockedError(AocdError):
+    """trying to access input before the unlock"""
+
+
 class PuzzleUnsolvedError(AocdError):
     """answer is unknown because user has not solved puzzle yet"""

aocd/get.py

@@ -11,16 +11,18 @@
 from logging import getLogger
 
 from .exceptions import AocdError
+from .exceptions import PuzzleLockedError
 from .models import default_user
 from .models import Puzzle
 from .models import User
 from .utils import AOC_TZ
+from .utils import blocker
 
 
 log = getLogger(__name__)
 
 
-def get_data(session=None, day=None, year=None):
+def get_data(session=None, day=None, year=None, block=False):
     """
     Get data for day (1-25) and year (>= 2015)
     User's session cookie is needed (puzzle inputs differ by user)
@@ -36,7 +38,14 @@ def get_data(session=None, day=None, year=None):
         year = most_recent_year()
         log.info("most recent year=%s", year)
     puzzle = Puzzle(year=year, day=day, user=user)
-    return puzzle.input_data
+    try:
+        return puzzle.input_data
+    except PuzzleLockedError:
+        if not block:
+            raise
+        q = block == "q"
+        blocker(quiet=q)
+        return puzzle.input_data
 
 
 def most_recent_year():

aocd/models.py

@@ -23,6 +23,7 @@
 
 from .exceptions import AocdError
 from .exceptions import PuzzleUnsolvedError
+from .exceptions import PuzzleLockedError
 from .utils import AOC_TZ
 from .version import __version__
 
@@ -158,6 +159,8 @@ def input_data(self):
             url=self.input_data_url, cookies=self.user.auth, headers=USER_AGENT
         )
         if not response.ok:
+            if response.status_code == 404:
+                raise PuzzleLockedError("{}/{:02d} not available yet".format(self.year, self.day))
             log.error("got %s status code token=%s", response.status_code, sanitized)
             log.error(response.text)
             raise AocdError("Unexpected response")

aocd/utils.py

@@ -1,13 +1,14 @@
 import sys
 import time
+import tzlocal
 from datetime import datetime
 from itertools import cycle
 from dateutil.tz import gettz
 
 AOC_TZ = gettz("America/New_York")
 
 
-def blocker(quiet=False, dt=0.1):
+def blocker(quiet=False, dt=0.1, datefmt="%-I:%M %p"):
     """
     This function just blocks until the next puzzle unlocks.
     Pass `quiet=True` to disable the spinner etc.
@@ -24,11 +25,13 @@ def blocker(quiet=False, dt=0.1):
         year += 1
     unlock = datetime(year, month, day, tzinfo=AOC_TZ)
     spinner = cycle(r"\|/-")
+    localzone = tzlocal.get_localzone()
     while datetime.now(tz=AOC_TZ) < unlock:
-        msg = "{} Unlock day {} at midnight EST ({:.0f}s remaining)"
+        local_unlock = unlock.astimezone(tz=localzone)
+        msg = "{} Unlock day {} at {:%s} ({:.0f}s remaining)" % datefmt
         remaining = (unlock - datetime.now(tz=AOC_TZ)).total_seconds()
         if not quiet:
-            sys.stdout.write(msg.format(next(spinner), unlock.day, remaining))
+            sys.stdout.write(msg.format(next(spinner), unlock.day, local_unlock, remaining))
             sys.stdout.flush()
         time.sleep(dt)
         if not quiet:

setup.py

@@ -36,6 +36,7 @@
         "beautifulsoup4",
         "pebble",
         'colorama; platform_system == "Windows"',
+        "tzlocal",
     ],
     options={"bdist_wheel": {"universal": "1"}},
 )

tests/test_get_data.py

@@ -4,6 +4,7 @@
 
 import aocd
 from aocd.exceptions import AocdError
+from aocd.exceptions import PuzzleLockedError
 
 
 def test_get_from_server(requests_mock):
@@ -45,19 +46,31 @@ def test_saved_data_is_reused_if_available(tmpdir, requests_mock):
 def test_server_error(requests_mock, caplog):
     mock = requests_mock.get(
         url="https://adventofcode.com/2101/day/1/input",
-        text="Not Found",
-        status_code=404,
+        text="AWS meltdown",
+        status_code=500,
     )
     with pytest.raises(AocdError("Unexpected response")):
         aocd.get_data(year=2101, day=1)
     assert mock.called
     assert mock.call_count == 1
     assert caplog.record_tuples == [
-        ("aocd.models", logging.ERROR, "got 404 status code token=...oken"),
-        ("aocd.models", logging.ERROR, "Not Found"),
+        ("aocd.models", logging.ERROR, "got 500 status code token=...oken"),
+        ("aocd.models", logging.ERROR, "AWS meltdown"),
     ]
 
 
+def test_puzzle_not_available_yet(requests_mock, caplog):
+    mock = requests_mock.get(
+        url="https://adventofcode.com/2101/day/1/input",
+        text="Not Found",
+        status_code=404,
+    )
+    with pytest.raises(PuzzleLockedError("2101/01 not available yet")):
+        aocd.get_data(year=2101, day=1)
+    assert mock.called
+    assert mock.call_count == 1
+
+
 def test_session_token_in_req_headers(requests_mock):
     mock = requests_mock.get("https://adventofcode.com/2018/day/1/input")
     aocd.get_data(year=2018, day=1)