Support fig.show() (#22)

* Support fig.show()

* isort

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

Author: harupy

docs/source/examples/index.rst

@@ -13,3 +13,4 @@ Examples
    fig-vars
    precode
    iframe-size
+   show

docs/source/examples/show.rst

@@ -0,0 +1,18 @@
+show
+====
+
+.. plotly::
+
+   import plotly.express as px
+
+   fig = px.scatter(x=[0, 1, 2, 3, 4], y=[0, 1, 4, 9, 16])
+   fig.show()
+
+
+.. plotly::
+   :fig-vars: figure
+
+   import plotly.express as px
+
+   figure = px.scatter(x=[0, 1, 2, 3, 4], y=[0, 1, 4, 9, 16])
+   figure.show()

sphinx_plotly_directive/__init__.py

@@ -172,7 +172,9 @@
     assign_last_line_into_variable,
     create_code_block,
     create_directive_block,
+    ends_with_show,
     save_plotly_figure,
+    strip_last_line,
 )
 
 
@@ -532,7 +534,11 @@ def run_code(code, code_path, ns=None, function_name=None, fig_vars=None):
 
         variable_name = "fig"
 
-        if function_name is not None:
+        if ends_with_show(code):
+            print(strip_last_line(code))
+            exec(strip_last_line(code), ns)
+            figs = [ns[fig_var] for fig_var in fig_vars] if fig_vars else [ns[variable_name]]
+        elif function_name is not None:
             exec(code, ns)
             exec(assign_last_line_into_variable(function_name + "()", variable_name), ns)
             figs = [ns[variable_name]]

sphinx_plotly_directive/utils.py

@@ -1,3 +1,4 @@
+import re
 import textwrap
 
 import plotly
@@ -132,3 +133,60 @@ def create_code_block(code, language=None):
             "",
         ]
     )
+
+
+def strip_last_line(code):
+    r"""
+    Strips the last line of the give code block
+
+    Parameters
+    ----------
+    code : str
+        Code to strip
+
+    Returns
+    -------
+    str:
+        Stripped code
+
+    Examples
+    --------
+    >>> strip_last_line("a")
+    ''
+    >>> strip_last_line("a\nb")
+    'a'
+    >>> strip_last_line("a\nb\nc")
+    'a\nb'
+    """
+    return "\n".join(code.strip().split("\n")[:-1])
+
+
+def ends_with_show(code):
+    r"""
+    Returns True if the last line of the given code block ends with `show()`
+
+    Parameters
+    ----------
+    code : str
+        Code that may contain a line that looks like `fig.show()`
+
+    Returns
+    -------
+    str:
+        Variable name of the object that calls `show()`
+
+    Examples
+    --------
+    >>> ends_with_show("fig.show()")  # simple
+    True
+    >>> ends_with_show("fig.show(1, a=2)")  # show with arguments
+    True
+    >>> ends_with_show("fig = dummy\nfig.show()\n")  # multiline
+    True
+    >>> ends_with_show("foo")  # doesn't contains `show`
+    False
+    """
+    # TODO: Use a more strict regular expression
+    pattern = r"^(.+)\.show\(.*\)$"
+    match = re.search(pattern, code.strip().split("\n")[-1], flags=re.DOTALL)
+    return bool(match)

tests/test_utils.py

@@ -1,12 +1,15 @@
 import os
 
 import plotly.express as px
+import pytest
 
 from sphinx_plotly_directive.utils import (
     assign_last_line_into_variable,
     create_code_block,
     create_directive_block,
+    ends_with_show,
     save_plotly_figure,
+    strip_last_line,
 )
 
 
@@ -63,3 +66,28 @@ def test_create_code_block():
 """.lstrip()
 
     assert create_code_block(code, "python") == expected
+
+
+@pytest.mark.parametrize(
+    "code, expected_result",
+    [
+        ("a", ""),
+        ("a\nb", "a"),
+        ("a\nb\nc", "a\nb"),
+    ],
+)
+def test_strip_last_line(code, expected_result):
+    assert strip_last_line(code) == expected_result
+
+
+@pytest.mark.parametrize(
+    "code, expected_result",
+    [
+        ("fig.show()", True),  # simple
+        ("fig.show(1, a=2)", True),  # contains arguments
+        ("fig = dummy\nfig.show()\n", True),  # multiline
+        ("fig", False),  # no show
+    ],
+)
+def test_ends_with_show(code, expected_result):
+    assert ends_with_show(code) == expected_result