Screenshot functionality

doc/builtins.rst

@@ -55,6 +55,15 @@ draw images to the screen ("blit" them).
         parameter. If ``image`` is a ``str`` then the named image will be
         loaded from the ``images/`` directory.
 
+    .. method:: screenshot()
+
+        Takes a screenshot of the entire game window and saves it to ``Pictures`` in your home directory.
+
+        Returns the path of the file that was written, as a string.
+
+        You can press F12 to take a screenshot at any time, without
+        needing to call this function.
+
     .. method:: draw.line(start, end, (r, g, b), width=1)
 
         Draw a line from start to end with a certain line width.

src/pgzero/game.py

@@ -3,6 +3,7 @@
 import time
 import types
 from time import perf_counter, sleep
+from traceback import print_exc
 
 import pygame
 import pgzero.clock
@@ -256,6 +257,14 @@ def key_down(event):
             if event.key == pygame.K_q and \
                     event.mod & (pygame.KMOD_CTRL | pygame.KMOD_META):
                 sys.exit(0)
+            # Default key for screenshots is F12.
+            if event.key == pygame.K_F12:
+                try:
+                    path = pgzero.screen.screen_instance.screenshot()
+                    print(f"Saved screenshot to {path}")
+                except Exception:
+                    print("ERROR while trying to take a screenshot with F12.")
+                    print_exc()
             self.keyboard._press(event.key)
             if user_key_down:
                 return user_key_down(event)

src/pgzero/runner.py

@@ -1,4 +1,5 @@
 from . import storage
+from . import screen
 from . import clock
 from . import loaders
 from . import __version__
@@ -217,6 +218,8 @@ def prepare_mod(mod):
 
     """
     storage.storage._set_filename_from_path(mod.__file__)
+    # Create the screen.screenshots instance.
+    screen._initialize_screenshots(mod.__file__)
     loaders.set_root(mod.__file__)
 
     # Copy pgzero builtins into system builtins

src/pgzero/screen.py

@@ -1,6 +1,11 @@
+import os
+import sys
+from datetime import datetime
+
 import numpy as np
 import pygame
 import pygame.draw
+import pygame.image
 
 from . import ptext
 from .rect import RECT_CLASSES, ZRect
@@ -25,6 +30,59 @@ def make_color(arg):
     return tuple(pygame.Color(arg))
 
 
+def _get_platform_screenshot_path():
+    r"""Get the screenshot directory for pgzero.
+    Under Windows, this is %USERPROFILE%\Pictures\pgzero.
+    Under Linux/MacOS, it's ~/Pictures/pgzero.
+    If the platform is unsupported, it defaults to the CWD."""
+    if sys.platform == "win32":
+        try:
+            home = os.environ["USERPROFILE"]
+        except KeyError:
+            raise KeyError("Couldn't find the user home directory for "
+                           "screenshots. Please set the %USERPROFILE% "
+                           "environment variable.")
+        return os.path.join(home, "Pictures", "pgzero")
+    elif sys.platform in ("linux", "linux2", "darwin"):
+        return os.path.expanduser(os.path.join("~", "Pictures", "pgzero"))
+    else:
+        print(f"WARNING: Device platform {sys.platform} not recognized, thus "
+              "no user folder found. Falling back to current directory to save"
+              " screenshots.", file=sys.stderr)
+        return os.path.join(os.getcwd(), "pgzero_screenshots")
+
+
+# This function is used to create the screenshot instance with the file name
+# given by runner.py but save it in the scope of screen.
+def _initialize_screenshots(file_path):
+    global screenshots
+    # Otherwise, create the instance of the Screenshots class used to
+    # take and save screenshots.
+    if not os.path.isabs(file_path):
+        file_path = os.path.abspath(file_path)
+    project_name, _ = os.path.splitext(os.path.basename(file_path))
+    screenshots = Screenshots(project_name)
+
+
+class Screenshots:
+    """Class to manage taking screenshots."""
+    def __init__(self, project_name):
+        self._project_name = project_name
+        self._path = _get_platform_screenshot_path()
+
+    def take(self, surface):
+        # Ensure that the directory for screenshots exists.
+        os.makedirs(self._path, exist_ok=True)
+
+        # Creates the filename, made up of the script name and a timestamp.
+        now = datetime.now()
+        filename = f"{self._project_name}-{now:%Y-%m-%d_%H:%M:%S}.png"
+        filepath = os.path.join(self._path, filename)
+        # Save the screenshot.
+        pygame.image.save(surface, filepath)
+        return filepath
+
+
 class SurfacePainter:
     """Interface to pygame.draw that is bound to a surface."""
 
@@ -162,6 +220,10 @@ def blit(self, image, pos):
             image = loaders.images.load(image)
         self.surface.blit(image, pos, None, pygame.BLEND_ALPHA_SDL2)
 
+    def screenshot(self):
+        """Takes a screenshot of the entire game window."""
+        return screenshots.take(self.surface)
+
     @property
     def draw(self):
         return SurfacePainter(self)

src/pgzero/storage.py

@@ -122,6 +122,10 @@ def save(self):
         """Save data to disk."""
         if not self and not self.loaded:
             return
+        # The save functionality seems to have been broken before this point,
+        # since saving manually in the game script did not make sure the save
+        # path actually existed. This fixes it.
+        Storage._ensure_save_path()
         try:
             data = json.dumps(self)
         except TypeError:

test/test_screen.py

@@ -1,5 +1,7 @@
 import sys
 import unittest
+from unittest.mock import patch
+from tempfile import TemporaryDirectory
 from pathlib import Path
 import os
 import warnings
@@ -9,7 +11,7 @@
 import pygame.image
 import pygame.surfarray
 
-from pgzero.screen import Screen
+import pgzero.screen as screen
 from pgzero.loaders import set_root, images
 from pgzero.rect import Rect, ZRect
 
@@ -96,7 +98,7 @@ def tearDownClass(cls):
         pygame.display.quit()
 
     def setUp(self):
-        self.screen = Screen()
+        self.screen = screen.Screen()
         self.screen._set_surface(self.surf)
         self.screen.clear()
 
@@ -238,6 +240,40 @@ def test_bounds(self):
             ZRect(0, 0, 200, 200)
         )
 
+    @patch("sys.platform", "win32")
+    @patch.dict("os.environ", {"USERPROFILE": r"c:\Users\user"})
+    def test_get_screenshot_path_windows(self):
+        r"""Screenshot path on Windows is %USERPROFILE%\Pictures\pgzero."""
+        result_path = screen._get_platform_screenshot_path()
+        self.assertEqual(result_path,
+                         os.path.join(r"c:\Users\user", "Pictures", "pgzero"))
+
+    @patch("sys.platform", "linux")
+    @patch.dict("os.environ", {"HOME": "/home/user"})
+    def test_get_screenshot_path_linux(self):
+        """Screenshot path on Linux or MacOS is ~/Pictures/pgzero."""
+        result_path = screen._get_platform_screenshot_path()
+        self.assertEqual(result_path,
+                         os.path.join("/home/user", "Pictures", "pgzero"))
+
+    @patch("sys.platform", "NOTHING")
+    def test_get_screenshot_path_other(self):
+        """If OS is not supported, CWD is used for screenshots."""
+        result_path = screen._get_platform_screenshot_path()
+        self.assertEqual(result_path,
+                         os.path.join(os.getcwd(), "pgzero_screenshots"))
+
+    @patch("sys.platform", "NOTHING")
+    def test_take_screenshot(self):
+        """Screenshot files are created and have the proper extension."""
+        with TemporaryDirectory("screenshot_testdir") as td:
+            os.chdir(td)
+            screen._initialize_screenshots(__file__)
+            self.screen.screenshot()
+            self.assertEqual(len(os.listdir("pgzero_screenshots")), 1)
+            ext = os.listdir("pgzero_screenshots")[0].split(".")[-1]
+            self.assertEqual(ext, "png")
+
 
 if __name__ == '__main__':
     unittest.main()