Add context managers for turtle.fill and turtle.poly

Co-authored-by: Marie Roald <[email protected]>
python · erlend-aasland · Jan 18, 2025 · Oct 6, 2024 · Oct 20, 2024 · Oct 27, 2024
commit 3379b3d6620a55a95712d8707f7ffc7b66fbea82
diff --git a/Lib/test/test_turtle.py b/Lib/test/test_turtle.py
@@ -4,6 +4,8 @@
 import unittest
 import unittest.mock
 import tempfile
+import sys
+from contextlib import contextmanager
 from test import support
 from test.support import import_helper
 from test.support import os_helper
@@ -527,6 +529,112 @@ def test_save(self) -> None:
                 assert f.read() == "postscript"
 
 
+class TestTurtle(unittest.TestCase):
+    @contextmanager
+    def patch_screen(self):
+        """Patch turtle._Screen for testing without a display.
+
+        We must patch the `_Screen` class itself instead of the `_Screen`
+        instance because instatiating it requires a display.
+        """
+        m = unittest.mock.MagicMock()
+        m.__class__ = turtle._Screen
+        m.mode.return_value = "standard"
+
+        patch = unittest.mock.patch('turtle._Screen.__new__', return_value=m)
+        try:
+            yield patch.__enter__()
+        finally:
+            patch.__exit__(*sys.exc_info())
+
+    def test_begin_end_fill(self):
+        """begin_fill and end_fill counter each other."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        self.assertFalse(t.filling())
+        t.begin_fill()
+        self.assertTrue(t.filling())
+        t.end_fill()
+        self.assertFalse(t.filling())
+
+    def test_fill(self):
+        """The context manager behaves like begin_ and end_ fill."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        self.assertFalse(t.filling())
+        with t.fill():
+            self.assertTrue(t.filling())
+        self.assertFalse(t.filling())
+
+    def test_fill_resets_after_exception(self):
+        """The context manager cleans up correctly after exceptions."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+        try:
+            with t.fill():
+                self.assertTrue(t.filling())
+                raise Exception
+        except Exception:
+            self.assertFalse(t.filling())
+
+    def test_fill_context_when_filling(self):
+        """The context manager works even when the turtle is already filling."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        t.begin_fill()
+        self.assertTrue(t.filling())
+        with t.fill():
+            self.assertTrue(t.filling())
+        self.assertFalse(t.filling())
+
+    def test_begin_end_poly(self):
+        """begin_fill and end_poly counter each other."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        self.assertFalse(t._creatingPoly)
+        t.begin_poly()
+        self.assertTrue(t._creatingPoly)
+        t.end_poly()
+        self.assertFalse(t._creatingPoly)
+
+    def test_poly(self):
+        """The context manager behaves like begin_ and end_ poly."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        self.assertFalse(t._creatingPoly)
+        with t.poly():
+            self.assertTrue(t._creatingPoly)
+        self.assertFalse(t._creatingPoly)
+
+    def test_poly_resets_after_exception(self):
+        """The context manager cleans up correctly after exceptions."""
+        with self.patch_screen():
+            t = turtle.Turtle()
+        try:
+            with t.poly():
+                self.assertTrue(t._creatingPoly)
+                raise Exception
+        except Exception:
+            self.assertFalse(t._creatingPoly)
+
+    def test_poly_context_when_creating_poly(self):
+        """The context manager works when the turtle is already creating poly.
+        """
+        with self.patch_screen():
+            t = turtle.Turtle()
+
+        t.begin_poly()
+        self.assertTrue(t._creatingPoly)
+        with t.poly():
+            self.assertTrue(t._creatingPoly)
+        self.assertFalse(t._creatingPoly)
+
+
 class TestModuleLevel(unittest.TestCase):
     def test_all_signatures(self):
         import inspect

diff --git a/Lib/turtle.py b/Lib/turtle.py
@@ -110,6 +110,8 @@
 from copy import deepcopy
 from tkinter import simpledialog
 
+from contextlib import contextmanager
+
 _tg_classes = ['ScrolledCanvas', 'TurtleScreen', 'Screen',
                'RawTurtle', 'Turtle', 'RawPen', 'Pen', 'Shape', 'Vec2D']
 _tg_screen_functions = ['addshape', 'bgcolor', 'bgpic', 'bye',
@@ -3382,6 +3384,14 @@ def filling(self):
         """
         return isinstance(self._fillpath, list)
 
+    @contextmanager
+    def fill(self):
+        self.begin_fill()
+        try:
+            yield
+        finally:
+            self.end_fill()
+
     def begin_fill(self):
         """Called just before drawing a shape to be filled.
 
@@ -3402,7 +3412,6 @@ def begin_fill(self):
             self.undobuffer.push(("beginfill", self._fillitem))
         self._update()
 
-
     def end_fill(self):
         """Fill the shape drawn after the call begin_fill().
 
@@ -3506,6 +3515,14 @@ def write(self, arg, move=False, align="left", font=("Arial", 8, "normal")):
         if self.undobuffer:
             self.undobuffer.cumulate = False
 
+    @contextmanager
+    def poly(self):
+        self.begin_poly()
+        try:
+            yield
+        finally:
+            self.end_poly()
+
     def begin_poly(self):
         """Start recording the vertices of a polygon.