Merge pull request #93 from wimglenn/stack

be more strict when crawling the stack

Author: wimglenn

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       matrix:
         os:
-        - ubuntu-latest
+        - ubuntu-20.04
         - macos-latest
         - windows-latest
         python-version:

README.rst

@@ -61,6 +61,12 @@ Install with pip
 
    pip install advent-of-code-data
 
+If you want to use this within a Jupyter notebook, there are some extra deps:
+
+.. code-block:: bash
+
+   pip install 'advent-of-code-data[nb]'
+
 **Puzzle inputs differ by user.**   So export your session ID, for example:
 
 .. code-block:: bash

aocd/__init__.py

@@ -17,6 +17,7 @@
 from . import transforms
 from . import utils
 from . import version
+from . import _ipykernel
 from .exceptions import AocdError
 from .exceptions import PuzzleUnsolvedError
 from .get import get_data
@@ -44,6 +45,7 @@
     "AocdError",
     "PuzzleUnsolvedError",
     "AOC_TZ",
+    "_ipykernel",
 ]
 __all__ += transforms.__all__
 

aocd/_ipykernel.py

@@ -0,0 +1,22 @@
+import os
+import re
+
+import requests
+
+
+def get_ipynb_path():
+    # inline imports to avoid hard-dependency on IPython/jupyter
+    import IPython
+    from jupyter_server import serverapp
+    from jupyter_server.utils import url_path_join
+    app = IPython.get_ipython().config["IPKernelApp"]
+    kernel_id = re.search(r"(?<=kernel-)[\w\-]+(?=\.json)", app["connection_file"])[0]
+    for serv in serverapp.list_running_servers():
+        url = url_path_join(serv["url"], "api/sessions")
+        resp = requests.get(url, params={"token": serv["token"]})
+        resp.raise_for_status()
+        for sess in resp.json():
+            if kernel_id == sess["kernel"]["id"]:
+                path = serv["root_dir"]
+                fname = sess["notebook"]["path"]
+                return os.path.join(path, fname)

aocd/get.py

@@ -17,6 +17,7 @@
 from .models import User
 from .utils import AOC_TZ
 from .utils import blocker
+from ._ipykernel import get_ipynb_path
 
 
 log = getLogger(__name__)
@@ -96,23 +97,36 @@ def get_day_and_year():
     """
     pattern_year = r"201[5-9]|202[0-9]"
     pattern_day = r"2[0-5]|1[0-9]|[1-9]"
-    stack = [f[0] for f in traceback.extract_stack()]
-    for name in stack:
-        basename = os.path.basename(name)
+    for frame in traceback.extract_stack():
+        filename = frame[0]
+        linetxt = frame[-1] or ""
+        basename = os.path.basename(filename)
         reasons_to_skip_frame = [
             not re.search(pattern_day, basename),  # no digits in filename
-            name == __file__,  # here
-            "importlib" in name,  # Python 3 import machinery
-            "/IPython/" in name,  # IPython adds a tonne of stack frames
-            name.startswith("<"),  # crap like <decorator-gen-57>
-            name.endswith("ython3"),  # ipython3 alias
+            filename == __file__,  # here
+            "importlib" in filename,  # Python 3 import machinery
+            "/IPython/" in filename,  # IPython adds a tonne of stack frames
+            filename.startswith("<"),  # crap like <decorator-gen-57>
+            filename.endswith("ython3"),  # ipython3 alias
             basename.startswith("pydev_ipython_console"),  # PyCharm Python Console
+            "aocd" not in linetxt,
+            "ipykernel" in filename,
         ]
         if not any(reasons_to_skip_frame):
-            log.debug("stack crawl found %s", name)
-            abspath = os.path.abspath(name)
+            log.debug("stack crawl found %s", filename)
+            abspath = os.path.abspath(filename)
             break
-        log.debug("skipping frame %s", name)
+        elif "ipykernel" in filename:
+            log.debug("stack crawl found %s, attempting to detect an .ipynb", filename)
+            try:
+                abspath = get_ipynb_path()
+            except Exception as err:
+                log.debug("failed getting .ipynb path with %s %s", type(err), err)
+            else:
+                if abspath and re.search(pattern_day, abspath):
+                    basename = os.path.basename(abspath)
+                    break
+        log.debug("skipping frame %s", filename)
     else:
         import __main__
         if getattr(__main__, "__file__", "<input>") == "<input>":
@@ -135,7 +149,7 @@ def get_day_and_year():
         assert not day.startswith("0"), "regex pattern_day must prevent any leading 0"
         day = int(day)
         assert 1 <= day <= 25, "regex pattern_day must only match numbers in range 1-25"
-        log.debug("year=%d day=%d", year, day)
+        log.debug("year=%s day=%s", year or "?", day)
         return day, year
     log.debug("giving up introspection for %s", abspath)
     raise AocdError("Failed introspection of day")

aocd/version.py

@@ -1 +1 @@
-__version__ = "1.2.3"
+__version__ = "1.3.0"

setup.py

@@ -45,4 +45,5 @@
         "tzlocal",
     ],
     options={"bdist_wheel": {"universal": "1"}},
+    extras_require={"nb": ["IPython", "jupyter-server"]}
 )

tests/test_aocd.py

@@ -9,7 +9,7 @@ def test_get_data_imported_from_class():
 
 
 def test_get_data_via_import(mocker):
-    fake_stack = [("~/2017/q23.py",)]
+    fake_stack = [("~/2017/q23.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     mock = mocker.patch("aocd._module.get_data", return_value="test data")
     from aocd import data
@@ -19,7 +19,7 @@ def test_get_data_via_import(mocker):
 
 
 def test_import_submit_binds_day_and_year(mocker):
-    fake_stack = [("~/2017/q23.py",)]
+    fake_stack = [("~/2017/q23.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     from aocd import submit
 
@@ -46,7 +46,7 @@ def test_get_data_via_import_in_interactive_mode(monkeypatch, mocker, freezer):
 
 
 def test_get_lines_via_import(mocker):
-    fake_stack = [("~/2017/q23.py",)]
+    fake_stack = [("~/2017/q23.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     mock = mocker.patch("aocd._module.get_data", return_value="line 1\nline 2\nline 3")
     from aocd import lines
@@ -56,7 +56,7 @@ def test_get_lines_via_import(mocker):
 
 
 def test_get_numbers_via_import(mocker):
-    fake_stack = [("~/2017/q23.py",)]
+    fake_stack = [("~/2017/q23.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     mock = mocker.patch("aocd._module.get_data", return_value="1\n2\n3")
     from aocd import numbers

tests/test_date_introspection.py

@@ -10,29 +10,54 @@ def test_get_day_and_year_fail_no_filename_on_stack():
 
 
 def test_get_day_and_year_from_stack(mocker):
-    fake_stack = [("xmas_problem_2016_25b_dawg.py",)]
+    fake_stack = [("xmas_problem_2016_25b_dawg.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     day, year = get_day_and_year()
     assert day == 25
     assert year == 2016
 
 
 def test_year_is_ambiguous(mocker):
-    fake_stack = [("~/2016/2017_q01.py",)]
+    fake_stack = [("~/2016/2017_q01.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     with pytest.raises(AocdError("Failed introspection of year")):
         get_day_and_year()
 
 
 def test_day_is_unknown(mocker):
-    fake_stack = [("~/2016.py",)]
+    fake_stack = [("~/2016.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     with pytest.raises(AocdError("Failed introspection of day")):
         get_day_and_year()
 
 
 def test_day_is_invalid(mocker):
-    fake_stack = [("~/2016/q27.py",)]
+    fake_stack = [("~/2016/q27.py", 1, "<test>", "from aocd import data")]
     mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
     with pytest.raises(AocdError("Failed introspection of day")):
         get_day_and_year()
+
+
+def test_ipynb_ok(mocker):
+    fake_stack = [("ipykernel/123456789.py", 1, "<test>", "from aocd import data")]
+    mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
+    mocker.patch("aocd.get.get_ipynb_path", return_value="puzzle-2020-03.py")
+    day, year = get_day_and_year()
+    assert day == 3
+    assert year == 2020
+
+
+def test_ipynb_fail(mocker):
+    fake_stack = [("ipykernel/123456789.py", 1, "<test>", "from aocd import data")]
+    mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
+    mocker.patch("aocd.get.get_ipynb_path", side_effect=ImportError)
+    with pytest.raises(AocdError("Failed introspection of filename")):
+        get_day_and_year()
+
+
+def test_ipynb_fail_no_numbers_in_ipynb_filename(mocker):
+    fake_stack = [("ipykernel/123456789.py", 1, "<test>", "from aocd import data")]
+    mocker.patch("aocd.get.traceback.extract_stack", return_value=fake_stack)
+    mocker.patch("aocd.get.get_ipynb_path", "puzzle.py")
+    with pytest.raises(AocdError("Failed introspection of filename")):
+        get_day_and_year()