diff --git a/doc/builtins.rst b/doc/builtins.rst index 4ac0e0fd..5957d70b 100644 --- a/doc/builtins.rst +++ b/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. diff --git a/src/pgzero/game.py b/src/pgzero/game.py index 4f242a85..30743c10 100644 --- a/src/pgzero/game.py +++ b/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) diff --git a/src/pgzero/runner.py b/src/pgzero/runner.py index e7c3891b..53d710a7 100644 --- a/src/pgzero/runner.py +++ b/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 diff --git a/src/pgzero/screen.py b/src/pgzero/screen.py index 08ac2321..32c124f2 100644 --- a/src/pgzero/screen.py +++ b/src/pgzero/screen.py @@ -1,5 +1,10 @@ +import os +import sys +from datetime import datetime + import pygame import pygame.draw +import pygame.image from . import ptext from .rect import RECT_CLASSES, ZRect @@ -24,6 +29,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.""" @@ -178,6 +236,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) diff --git a/src/pgzero/storage.py b/src/pgzero/storage.py index 652c64fc..082dde4e 100644 --- a/src/pgzero/storage.py +++ b/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: diff --git a/test/test_screen.py b/test/test_screen.py index 1246a81b..e20798be 100644 --- a/test/test_screen.py +++ b/test/test_screen.py @@ -1,7 +1,11 @@ import sys import unittest +from unittest.mock import patch +from tempfile import TemporaryDirectory from pathlib import Path import os +import ntpath +import posixpath import warnings import numpy as np @@ -9,7 +13,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 +100,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 +242,42 @@ def test_bounds(self): ZRect(0, 0, 200, 200) ) + @patch("sys.platform", "win32") + @patch("os.path", ntpath) + @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("os.path", posixpath) + @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()