Merge pull request #10 from gambitproject/jupyter-bits

Simplify API

Author: edwardchalstrey1

.github/workflows/test.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.13"]
+        python-version: ["3.10", "3.13"]
 
     steps:
     - uses: actions/checkout@v4

README.md

@@ -2,6 +2,9 @@
 
 Game tree drawing tool for extensive form games that generates TikZ code, LaTeX documents, PDFs, and PNGs.
 
+Pass in an extensive form game file in `.ef` format with layout formatting, and `draw_tree` will generate a visual representation of the game tree.
+You can also pass in a file in `.efg` format, which will be converted to `.ef` internally, applying a default layout.
+
 ## Installation
 
 Clone the repo and install the package using pip:
@@ -14,7 +17,7 @@ pip install -e .
 
 ## Requirements
 
-- Python 3.9+ (tested on 3.13)
+- Python 3.10+ (tested on 3.13)
 - LaTeX with TikZ (for PDF/PNG generation)
 - (optional) ImageMagick or Ghostscript or Poppler (for PNG generation)
 
@@ -73,18 +76,11 @@ generate_png('games/example.ef', output_png='mygame.png', scale_factor=0.8)    #
 
 ### Rendering in Jupyter Notebooks
 
-First install the requirements, which include the `jupyter-tikz` extension:
-```bash
-pip install -r requirements.txt
-```
-
 In a Jupyter notebook, run:
 
 ```python
-%load_ext jupyter_tikz
 from draw_tree import draw_tree
-example_tikz = draw_tree('games/example.ef')
-get_ipython().run_cell_magic("tikz", "", example_tikz)
+draw_tree('games/example.ef')
 ```
 
 ## Developer docs: Testing

pyproject.toml

@@ -25,8 +25,10 @@ classifiers = [
 ]
 keywords = ["game theory", "tikz", "visualization", "trees", "economics"]
 
+# Required runtime dependencies (previously optional under the 'jupyter' extra)
+dependencies = ["jupyter-tikz", "ipykernel"]
+
 [project.optional-dependencies]
-jupyter = ["jupyter-tikz", "ipykernel"]
 dev = ["pytest>=7.0.0", "pytest-cov"]
 
 [project.scripts]

requirements.txt

@@ -9,6 +9,7 @@
 
 from .core import (
     draw_tree,
+    generate_tikz,
     generate_tex,
     generate_pdf,
     generate_png,
@@ -19,6 +20,7 @@
 
 __all__ = [
     "draw_tree",
+    "generate_tikz",
     "generate_tex", 
     "generate_pdf",
     "generate_png",

src/draw_tree/__init__.py

@@ -15,6 +15,7 @@
 
 from pathlib import Path
 from typing import List, Optional 
+from IPython.core.getipython import get_ipython
 
 # Constants
 DEFAULTFILE: str = "example.ef"
@@ -1250,7 +1251,7 @@ def ef_to_tex(ef_file: str, scale_factor: float = 1.0, show_grid: bool = False)
         scale = original_scale
         grid = original_grid
 
-def draw_tree(ef_file: str, scale_factor: float = 1.0, show_grid: bool = False) -> str:
+def generate_tikz(ef_file: str, scale_factor: float = 1.0, show_grid: bool = False) -> str:
     """
     Generate complete TikZ code from an extensive form (.ef) file.
     
@@ -1323,6 +1324,39 @@ def draw_tree(ef_file: str, scale_factor: float = 1.0, show_grid: bool = False)
     return tikz_code
 
 
+def draw_tree(ef_file: str, scale_factor: float = 1.0, show_grid: bool = False) -> Optional[str]:
+    """
+    Generate TikZ code and display in Jupyter notebooks.
+    
+    Args:
+        ef_file: Path to the .ef file to process.
+        scale_factor: Scale factor for the diagram (default: 1.0).
+        show_grid: Whether to show grid lines (default: False).
+        
+    Returns:
+        The result of the Jupyter cell magic execution, or the TikZ code string
+        if cell magic fails.
+    """
+    # Ensure we are in a Jupyter notebook environment
+    ip = get_ipython()
+    if ip:
+        # Only attempt to load the extension if it's not already loaded
+        em = getattr(ip, 'extension_manager', None)
+        loaded = getattr(em, 'loaded', None)
+        try:
+            jpt_loaded = 'jupyter_tikz' in loaded  # type: ignore
+        except Exception:
+            jpt_loaded = False
+        if not jpt_loaded:
+            ip.run_line_magic("load_ext", "jupyter_tikz")
+
+        # Generate TikZ code and execute cell magic
+        tikz_code = generate_tikz(ef_file, scale_factor, show_grid)
+        return ip.run_cell_magic("tikz", "", tikz_code)
+    else:
+        raise EnvironmentError("draw_tree function requires a Jupyter notebook environment.")
+
+
 def latex_wrapper(tikz_code: str) -> str:
     """
     Wrap TikZ code in a complete LaTeX document.
@@ -1391,8 +1425,8 @@ def generate_tex(ef_file: str, output_tex: Optional[str] = None, scale_factor: f
         except Exception:
             pass
 
-    # Generate TikZ content using draw_tree
-    tikz_content = draw_tree(ef_file, scale_factor, show_grid)
+    # Generate TikZ content using generate_tikz
+    tikz_content = generate_tikz(ef_file, scale_factor, show_grid)
 
     # Wrap in complete LaTeX document
     latex_document = latex_wrapper(tikz_content)
@@ -1430,8 +1464,8 @@ def generate_pdf(ef_file: str, output_pdf: Optional[str] = None, scale_factor: f
         ef_path = Path(ef_file)
         output_pdf = ef_path.with_suffix('.pdf').name
 
-    # Generate TikZ content using draw_tree
-    tikz_content = draw_tree(ef_file, scale_factor, show_grid)
+    # Generate TikZ content using generate_tikz
+    tikz_content = generate_tikz(ef_file, scale_factor, show_grid)
 
     # Create LaTeX wrapper document
     latex_document = latex_wrapper(tikz_content)
@@ -2233,7 +2267,7 @@ def emit_node(n: 'DefaultLayout.Node'):
 
 
 def efg_to_ef(efg_file: str) -> str:
-    """Convert a Gambit .efg file to the `.ef` format used by draw_tree.
+    """Convert a Gambit .efg file to the `.ef` format used by generate_tikz.
 
     The function implements a focused parser and deterministic layout
     heuristics for producing `.ef` directives from a conservative subset of

src/draw_tree/core.py

@@ -263,7 +263,7 @@ def test_draw_tree_basic(self):
             ef_file_path = ef_file.name
 
         try:
-            result = draw_tree.draw_tree(ef_file_path)
+            result = draw_tree.generate_tikz(ef_file_path)
 
             # Verify the result contains expected components
             assert isinstance(result, str)
@@ -282,6 +282,69 @@ def test_draw_tree_basic(self):
         finally:
             os.unlink(ef_file_path)
 
+    def test_draw_tree_raises_when_no_ipython(self):
+        """When IPython is not available, draw_tree should raise EnvironmentError."""
+        with patch('draw_tree.core.get_ipython', return_value=None):
+            with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.ef') as ef_file:
+                ef_file.write("player 1\n")
+                ef_file.write("level 0 node root player 1\n")
+                ef_file_path = ef_file.name
+            try:
+                with pytest.raises(EnvironmentError):
+                    draw_tree.draw_tree(ef_file_path)
+            finally:
+                os.unlink(ef_file_path)
+
+    def test_draw_tree_calls_ipython_magic_when_available(self):
+        """When IPython is available, draw_tree should load the jupyter_tikz
+        extension if needed and call the tikz cell magic with the generated code.
+        """
+        # Create a simple .ef file for testing
+        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.ef') as ef_file:
+            ef_file.write("player 1\n")
+            ef_file.write("level 0 node root player 1\n")
+            ef_file_path = ef_file.name
+
+        class DummyEM:
+            def __init__(self, loaded=None):
+                self.loaded = loaded or set()
+
+        class DummyIP:
+            def __init__(self, em):
+                self.extension_manager = em
+                self._loaded_magics = []
+                self._run_cell_magic_calls = []
+
+            def run_line_magic(self, name, arg):
+                # record that load_ext was called
+                self._loaded_magics.append((name, arg))
+
+            def run_cell_magic(self, magic_name, args, code):
+                # record call and return a sentinel
+                self._run_cell_magic_calls.append((magic_name, args, code))
+                return "MAGIC-RESULT"
+
+        try:
+            # Case 1: extension already loaded
+            em = DummyEM(loaded={'jupyter_tikz'})
+            ip = DummyIP(em)
+            with patch('draw_tree.core.get_ipython', return_value=ip):
+                res = draw_tree.draw_tree(ef_file_path)
+                # Should call run_cell_magic and return its value
+                assert res == "MAGIC-RESULT"
+
+            # Case 2: extension not loaded -> run_line_magic should be called
+            em2 = DummyEM(loaded=set())
+            ip2 = DummyIP(em2)
+            with patch('draw_tree.core.get_ipython', return_value=ip2):
+                res2 = draw_tree.draw_tree(ef_file_path)
+                assert res2 == "MAGIC-RESULT"
+                # run_line_magic should have been called to load the extension
+                assert ('load_ext', 'jupyter_tikz') in ip2._loaded_magics
+
+        finally:
+            os.unlink(ef_file_path)
+
     def test_draw_tree_with_options(self):
         """Test draw_tree with different options."""
         # Create a simple .ef file for testing
@@ -292,15 +355,15 @@ def test_draw_tree_with_options(self):
 
         try:
             # Test with scale
-            result_scaled = draw_tree.draw_tree(ef_file_path, scale_factor=2.0)
+            result_scaled = draw_tree.generate_tikz(ef_file_path, scale_factor=2.0)
             assert "scale=2" in result_scaled
 
             # Test with grid
-            result_grid = draw_tree.draw_tree(ef_file_path, show_grid=True)
+            result_grid = draw_tree.generate_tikz(ef_file_path, show_grid=True)
             assert "\\draw [help lines, color=green]" in result_grid
 
             # Test without grid (default)
-            result_no_grid = draw_tree.draw_tree(ef_file_path, show_grid=False)
+            result_no_grid = draw_tree.generate_tikz(ef_file_path, show_grid=False)
             assert "% \\draw [help lines, color=green]" in result_no_grid
 
         finally:
@@ -310,15 +373,15 @@ def test_draw_tree_missing_files(self):
         """Test draw_tree with missing files."""
         # Test with missing .ef file
         with pytest.raises(FileNotFoundError):
-            draw_tree.draw_tree("nonexistent.ef")
+            draw_tree.generate_tikz("nonexistent.ef")
 
         # Test with valid .ef file (should work with built-in macros)
         with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.ef') as ef_file:
             ef_file.write("player 1\nlevel 0 node root player 1\n")
             ef_file_path = ef_file.name
 
         try:
-            result = draw_tree.draw_tree(ef_file_path)
+            result = draw_tree.generate_tikz(ef_file_path)
             # Should work with built-in macros
             assert "\\begin{tikzpicture}" in result
             assert "\\newcommand\\chancecolor{red}" in result