From c1f337f50b160f5ed009e5039e94a4c5f10df749 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Thu, 13 Feb 2025 20:07:40 -0500 Subject: [PATCH] docs: add check code output for result="ansi" (#3435) * Add ast parsing to determine whether we should include result="ansi". It's not meant to be perfect, but will hopefully catch the most common cases. Still requires manual review. * Ideally we could suppress output in markdown-exec in the future. --- docs/_scripts/notebook_convert.py | 96 +++++++++++++++++++ .../md_executable/index.md.j2 | 2 +- .../unit_tests/test_notebook_conversion.py | 33 ++++++- 3 files changed, 128 insertions(+), 3 deletions(-) diff --git a/docs/_scripts/notebook_convert.py b/docs/_scripts/notebook_convert.py index a23016c2a..66a56ff88 100644 --- a/docs/_scripts/notebook_convert.py +++ b/docs/_scripts/notebook_convert.py @@ -64,6 +64,100 @@ def _rewrite_cell_magic(code: str) -> str: return "\n".join(rewritten_lines) +class PrintCallVisitor(ast.NodeVisitor): + """ + This visitor sets self.has_print to True if it encounters a call + to a print within the global scope. + + This should catch calls to print(), print_stream(), etc. (Prefixed with "print"). + + May have some false positives, but it's not meant to be perfect. + + Temporary code for notebook conversion. + """ + + def __init__(self): + self.has_print = False + self.scope_level = 0 # counter to track whether we're inside a def/lambda + + def visit_FunctionDef(self, node): + self.scope_level += 1 + self.generic_visit(node) + self.scope_level -= 1 + + def visit_AsyncFunctionDef(self, node): + self.scope_level += 1 + self.generic_visit(node) + self.scope_level -= 1 + + def visit_Lambda(self, node): + self.scope_level += 1 + self.generic_visit(node) + self.scope_level -= 1 + + def visit_ClassDef(self, node): + self.scope_level += 1 + self.generic_visit(node) + self.scope_level -= 1 + + def visit_Call(self, node): + # Only consider calls when not inside a function definition. + if self.scope_level == 0: + if isinstance(node.func, ast.Name) and node.func.id.startswith("print"): + self.has_print = True + self.generic_visit(node) + + +def _has_output(source: str) -> bool: + """Determine if the code block is expected to produce output. + + Args: + source (str): The source code of the code block. + + Returns: + True if the code block is expected to produce output, False otherwise. + + Must meet the following conditions: + + 1. There is a call to a printing function (name starts with "print") + that is not inside a function definition. + 2. The last top-level statement is an expression that is valid if: + - It is any expression (including calls) AND + - It is NOT a call to `display(...)`. + + `display` isn't handled currently by markdown-exec + """ + try: + tree = ast.parse(source) + except SyntaxError: + return False + + # Condition (1): Check for a global print-like call. + visitor = PrintCallVisitor() + visitor.visit(tree) + condition_a = visitor.has_print + + # Condition (2): Check the last top-level statement. + condition_b = False + if tree.body: + last_stmt = tree.body[-1] + if isinstance(last_stmt, ast.Expr): + # If the expression is a call, ensure it's not a call to "display" + if isinstance(last_stmt.value, ast.Call): + if ( + isinstance(last_stmt.value.func, ast.Name) + and last_stmt.value.func.id == "display" + ): + condition_b = False # exclude display-wrapped expressions + else: + condition_b = True + else: + # Any other expression qualifies. + condition_b = True + + return condition_a or condition_b + + def _convert_links_in_markdown(markdown: str) -> str: """Convert links present in notebook markdown cells to standardized format. @@ -161,6 +255,8 @@ class EscapePreprocessor(Preprocessor): cell.source = _rewrite_cell_magic(source) cell.metadata["language"] = "shell" + cell.metadata["has_output"] = _has_output(source) + # Remove noqa comments cell.source = re.sub(r"#\s*noqa.*$", "", cell.source, flags=re.MULTILINE) # escape ``` in code diff --git a/docs/_scripts/notebook_convert_templates/md_executable/index.md.j2 b/docs/_scripts/notebook_convert_templates/md_executable/index.md.j2 index 25f18fc97..58e687fb4 100644 --- a/docs/_scripts/notebook_convert_templates/md_executable/index.md.j2 +++ b/docs/_scripts/notebook_convert_templates/md_executable/index.md.j2 @@ -8,7 +8,7 @@ {%- elif cell.metadata.get('language') == "shell" -%} shell {%- elif 'name' in nb.metadata.get('language_info', {}) -%} - {{ nb.metadata.language_info.name }}{% if cell.metadata.exec|default(false) %} exec="on" source="above" session="1" result="ansi"{% endif %} + {{ nb.metadata.language_info.name }}{% if cell.metadata.exec|default(false) %} exec="on" source="above" session="1"{% if cell.metadata.has_output|default(false) %} result="ansi"{% endif %}{% endif %} {%- endif %} {{ cell.source}} ``` diff --git a/docs/tests/unit_tests/test_notebook_conversion.py b/docs/tests/unit_tests/test_notebook_conversion.py index ffe7a21a0..61740b37d 100644 --- a/docs/tests/unit_tests/test_notebook_conversion.py +++ b/docs/tests/unit_tests/test_notebook_conversion.py @@ -1,8 +1,11 @@ import nbformat - -from _scripts.notebook_convert import md_executable, _convert_links_in_markdown import pytest +from _scripts.notebook_convert import ( + _convert_links_in_markdown, + md_executable, + _has_output, +) EXPECTED_OUTPUT = """\ ```python exec="on" source="above" session="1" result="ansi" @@ -59,6 +62,32 @@ def test_convert_input_cell() -> None: assert markdown == STDIN_OUTPUT +NO_STDOUT_EXPECTED = """\ +```python exec="on" source="above" session="1" +display(x) +``` +""" + + +def test_convert_block_without_output() -> None: + notebook = nbformat.v4.new_notebook() + notebook.metadata.language_info = {"name": "python", "version": "3.11"} + notebook.cells.append(nbformat.v4.new_code_cell("display(x)")) + markdown, _ = md_executable.from_notebook_node(notebook) + assert markdown == NO_STDOUT_EXPECTED + + +def test_has_output() -> None: + """Test if a given code block is expected to have output.""" + assert _has_output("print('Hello, world!')") is True + assert _has_output("print_stream(some_iterable)") is True + assert _has_output("foo.y") is True + assert _has_output("display(x)") is False + assert _has_output("assert 1 == 1") is False + assert _has_output("def foo(): pass") is False + assert _has_output("import foobar") is False + + @pytest.mark.parametrize( "source, expected", [