From 437f0b250c248addc5bca926d40d3ea119734d22 Mon Sep 17 00:00:00 2001
From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
Date: Tue, 7 Jan 2025 11:01:53 +0000
Subject: [PATCH] simplify

---
 utils/remove_docstring_examples.py | 47 +++++++++++++-----------------
 1 file changed, 21 insertions(+), 26 deletions(-)

diff --git a/utils/remove_docstring_examples.py b/utils/remove_docstring_examples.py
index 25dcc9059..c55a64387 100644
--- a/utils/remove_docstring_examples.py
+++ b/utils/remove_docstring_examples.py
@@ -5,41 +5,36 @@
 from ast import NodeVisitor
 
 
+def process_node(node: ast.FunctionDef | ast.ClassDef) -> tuple[int, int] | None:
+    """If node contains a docstring example, return start and end lines."""
+    if (
+        node.body
+        and isinstance(expr := node.body[0], ast.Expr)
+        and isinstance(value := expr.value, ast.Constant)
+        and isinstance(docstring := value.value, str)
+        and "Examples:" in docstring
+        and value.end_lineno is not None
+    ):
+        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.
+        return (examples_line_start - 1, value.end_lineno - 1)
+    return None
+
+
 class Visitor(NodeVisitor):
     def __init__(self, file: str) -> None:
         self.file = file
         self.to_remove: list[tuple[int, int]] = []
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: N802
-        if (
-            node.body
-            and isinstance(expr := node.body[0], ast.Expr)
-            and isinstance(value := expr.value, ast.Constant)
-            and isinstance(docstring := value.value, str)
-            and "Examples:" in docstring
-            and value.end_lineno is not None
-        ):
-            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 - 1, value.end_lineno - 1))
-
+        if removal := process_node(node):
+            self.to_remove.append(removal)
         self.generic_visit(node)
 
     def visit_ClassDef(self, node: ast.ClassDef) -> None:  # noqa: N802
-        if (
-            node.body
-            and isinstance(expr := node.body[0], ast.Expr)
-            and isinstance(value := expr.value, ast.Constant)
-            and isinstance(docstring := value.value, str)
-            and "Examples:" in docstring
-            and value.end_lineno is not None
-        ):
-            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))
-
+        if removal := process_node(node):
+            self.to_remove.append(removal)
         self.generic_visit(node)