add test

tests/utils_test.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import string
+import subprocess
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import hypothesis.strategies as st
@@ -271,3 +273,36 @@ def test_generate_temporary_column_name_raise() -> None:
         match="Internal Error: Narwhals was not able to generate a column name with ",
     ):
         nw.generate_temporary_column_name(n_bytes=1, columns=columns)
+
+
+def test_remove_docstring_examples(tmp_path: pytest.TempPathFactory) -> None:
+    src = (
+        "def foo(bar):\n"
+        '    """This is a docstring.\n'
+        "\n"
+        "    Arguments:\n"
+        "        bar\n"
+        "\n"
+        "    Examples:\n"
+        "        >>> foo(3)\n"
+        '    """\n'
+        "    print(bar)\n"
+    )
+    file = tmp_path / "t.py"  # type: ignore[operator]
+    file.write_text(src)
+    subprocess.run(  # noqa: S603
+        ["python", str(Path("utils") / "remove_docstring_examples.py"), str(file)],  # noqa: S607
+        check=False,
+    )
+    result = file.read_text()
+    expected = (
+        "def foo(bar):\n"
+        '    """This is a docstring.\n'
+        "\n"
+        "    Arguments:\n"
+        "        bar\n"
+        "\n"
+        '    """\n'
+        "    print(bar)\n"
+    )
+    assert result == expected

utils/remove_docstring_examples.py

@@ -21,8 +21,8 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: N802
         ):
             docstring_lines = [line.strip() for line in docstring.splitlines()]
             examples_line_start = value.lineno + docstring_lines.index("Examples:")
-            examples_line_end = value.end_lineno
-            self.to_remove.append((examples_line_start, examples_line_end))
+            # lineno is 1-indexed so we subtract 1.
+            self.to_remove.append((examples_line_start - 1, value.end_lineno - 1))
 
         self.generic_visit(node)
 
@@ -35,10 +35,10 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:  # noqa: N802
             and "Examples:" in docstring
             and value.end_lineno is not None
         ):
-            stripped_lines = [line.strip() for line in docstring.splitlines()]
-            examples_start = value.lineno + stripped_lines.index("Examples:")
-            # lineno is 1-indexed, so we subtract 1 to have 0-indexed numbers.
-            self.to_remove.append((examples_start - 1, value.end_lineno - 1))
+            docstring_lines = [line.strip() for line in docstring.splitlines()]
+            examples_line_start = value.lineno + docstring_lines.index("Examples:")
+            # lineno is 1-indexed so we subtract 1.
+            self.to_remove.append((examples_line_start, value.end_lineno))
 
         self.generic_visit(node)
 
@@ -59,4 +59,4 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:  # noqa: N802
             for examples_start, examples_end in removals:
                 del lines[examples_start:examples_end]
             with open(file, "w") as fd:
-                fd.write("\n".join(lines))
+                fd.write("\n".join(lines) + "\n")