Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 5 additions & 9 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
4 changes: 0 additions & 4 deletions requirements.txt

This file was deleted.

2 changes: 2 additions & 0 deletions src/draw_tree/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

from .core import (
draw_tree,
generate_tikz,
generate_tex,
generate_pdf,
generate_png,
Expand All @@ -19,6 +20,7 @@

__all__ = [
"draw_tree",
"generate_tikz",
"generate_tex",
"generate_pdf",
"generate_png",
Expand Down
46 changes: 40 additions & 6 deletions src/draw_tree/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
75 changes: 69 additions & 6 deletions tests/test_drawtree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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
Expand Down
Loading