gh-123614: Add save function to turtle.py (#123617)

diff --git a/Doc/library/turtle.rst b/Doc/library/turtle.rst
index afda3685d606bb74e63840ef86130e49446e76b6..da801d4dc1f5b3672061ab3c4f914cf4a9c7277a 100644 (file)
--- a/Doc/library/turtle.rst
+++ b/Doc/library/turtle.rst
@@ -427,6 +427,7 @@ Input methods
 Methods specific to Screen
    | :func:`bye`
    | :func:`exitonclick`
+   | :func:`save`
    | :func:`setup`
    | :func:`title`
 
@@ -2269,6 +2270,24 @@ Methods specific to Screen, not inherited from TurtleScreen
    client script.
 
 
+.. function:: save(filename, overwrite=False)
+
+   Save the current turtle drawing (and turtles) as a PostScript file.
+
+   :param filename: the path of the saved PostScript file
+   :param overwrite: if ``False`` and there already exists a file with the given
+                     filename, then the function will raise a
+                     ``FileExistsError``. If it is ``True``, the file will be
+                     overwritten.
+
+   .. doctest::
+      :skipif: _tkinter is None
+
+      >>> screen.save("my_drawing.ps")
+      >>> screen.save("my_drawing.ps", overwrite=True)
+
+   .. versionadded:: 3.14
+
 .. function:: setup(width=_CFG["width"], height=_CFG["height"], startx=_CFG["leftright"], starty=_CFG["topbottom"])
 
    Set the size and position of the main window.  Default values of arguments

diff --git a/Lib/test/test_turtle.py b/Lib/test/test_turtle.py
index 14121a590a50267ca42ca3ec57e21c6a15a6334a..c75a002a89b4c40a44cecf0dca0b6f23a34aab81 100644 (file)
--- a/Lib/test/test_turtle.py
+++ b/Lib/test/test_turtle.py
@@ -1,5 +1,9 @@
+import os
 import pickle
+import re
 import unittest
+import unittest.mock
+import tempfile
 from test import support
 from test.support import import_helper
 from test.support import os_helper
@@ -130,6 +134,7 @@ class VectorComparisonMixin:
             self.assertAlmostEqual(
                 i, j, msg='values at index {} do not match'.format(idx))
 
+
 class Multiplier:
 
     def __mul__(self, other):
@@ -461,6 +466,67 @@ class TestTPen(unittest.TestCase):
             self.assertTrue(tpen.isdown())
 
 
+class TestTurtleScreen(unittest.TestCase):
+    def test_save_raises_if_wrong_extension(self) -> None:
+        screen = unittest.mock.Mock()
+
+        msg = "Unknown file extension: '.png', must be one of {'.ps', '.eps'}"
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            self.assertRaisesRegex(ValueError, re.escape(msg))
+        ):
+            turtle.TurtleScreen.save(screen, os.path.join(tmpdir, "file.png"))
+
+    def test_save_raises_if_parent_not_found(self) -> None:
+        screen = unittest.mock.Mock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            parent = os.path.join(tmpdir, "unknown_parent")
+            msg = f"The directory '{parent}' does not exist. Cannot save to it"
+
+            with self.assertRaisesRegex(FileNotFoundError, re.escape(msg)):
+                turtle.TurtleScreen.save(screen, os.path.join(parent, "a.ps"))
+
+    def test_save_raises_if_file_found(self) -> None:
+        screen = unittest.mock.Mock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            file_path = os.path.join(tmpdir, "some_file.ps")
+            with open(file_path, "w") as f:
+                f.write("some text")
+
+            msg = (
+                f"The file '{file_path}' already exists. To overwrite it use"
+                " the 'overwrite=True' argument of the save function."
+            )
+            with self.assertRaisesRegex(FileExistsError, re.escape(msg)):
+                turtle.TurtleScreen.save(screen, file_path)
+
+    def test_save_overwrites_if_specified(self) -> None:
+        screen = unittest.mock.Mock()
+        screen.cv.postscript.return_value = "postscript"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            file_path = os.path.join(tmpdir, "some_file.ps")
+            with open(file_path, "w") as f:
+                f.write("some text")
+
+            turtle.TurtleScreen.save(screen, file_path, overwrite=True)
+            with open(file_path) as f:
+                assert f.read() == "postscript"
+
+    def test_save(self) -> None:
+        screen = unittest.mock.Mock()
+        screen.cv.postscript.return_value = "postscript"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            file_path = os.path.join(tmpdir, "some_file.ps")
+
+            turtle.TurtleScreen.save(screen, file_path)
+            with open(file_path) as f:
+                assert f.read() == "postscript"
+
+
 class TestModuleLevel(unittest.TestCase):
     def test_all_signatures(self):
         import inspect

diff --git a/Lib/turtle.py b/Lib/turtle.py
index 99850ae5efe348a6e3ce1f23aca7016c0288b918..8a5801f2efe625aaf85d87beff53836ccb5e017c 100644 (file)
--- a/Lib/turtle.py
+++ b/Lib/turtle.py
@@ -106,6 +106,7 @@ import inspect
 import sys
 
 from os.path import isfile, split, join
+from pathlib import Path
 from copy import deepcopy
 from tkinter import simpledialog
 
@@ -115,7 +116,7 @@ _tg_screen_functions = ['addshape', 'bgcolor', 'bgpic', 'bye',
         'clearscreen', 'colormode', 'delay', 'exitonclick', 'getcanvas',
         'getshapes', 'listen', 'mainloop', 'mode', 'numinput',
         'onkey', 'onkeypress', 'onkeyrelease', 'onscreenclick', 'ontimer',
-        'register_shape', 'resetscreen', 'screensize', 'setup',
+        'register_shape', 'resetscreen', 'screensize', 'save', 'setup',
         'setworldcoordinates', 'textinput', 'title', 'tracer', 'turtles', 'update',
         'window_height', 'window_width']
 _tg_turtle_functions = ['back', 'backward', 'begin_fill', 'begin_poly', 'bk',
@@ -1492,6 +1493,39 @@ class TurtleScreen(TurtleScreenBase):
         """
         return self._resize(canvwidth, canvheight, bg)
 
+    def save(self, filename, *, overwrite=False):
+        """Save the drawing as a PostScript file
+
+        Arguments:
+        filename -- a string, the path of the created file.
+                    Must end with '.ps' or '.eps'.
+
+        Optional arguments:
+        overwrite -- boolean, if true, then existing files will be overwritten
+
+        Example (for a TurtleScreen instance named screen):
+        >>> screen.save('my_drawing.eps')
+        """
+        filename = Path(filename)
+        if not filename.parent.exists():
+            raise FileNotFoundError(
+                f"The directory '{filename.parent}' does not exist."
+                " Cannot save to it."
+            )
+        if not overwrite and filename.exists():
+            raise FileExistsError(
+                f"The file '{filename}' already exists. To overwrite it use"
+                " the 'overwrite=True' argument of the save function."
+            )
+        if (ext := filename.suffix) not in {".ps", ".eps"}:
+            raise ValueError(
+                f"Unknown file extension: '{ext}',"
+                 " must be one of {'.ps', '.eps'}"
+            )
+
+        postscript = self.cv.postscript()
+        filename.write_text(postscript)
+
     onscreenclick = onclick
     resetscreen = reset
     clearscreen = clear

diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-09-02-20-39-10.gh-issue-123614.26TMHp.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-02-20-39-10.gh-issue-123614.26TMHp.rst
new file mode 100644 (file)
index 0000000..64a5eac
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-02-20-39-10.gh-issue-123614.26TMHp.rst
@@ -0,0 +1,2 @@
+Add :func:`turtle.save` to easily save Turtle drawings as PostScript files.
+Patch by Marie Roald and Yngve Mardal Moe.