Allow PyGame2 and backport revised image related tests, so they now p…

…ass.

setup.py

@@ -8,7 +8,7 @@
     LONG_DESCRIPTION = f.read()
 
 install_requires = [
-    'pygame>=1.9.2, <2.0',
+    'pygame>=1.9.2',
     'numpy',
 ]
 

test/test_screen.py

@@ -1,6 +1,7 @@
 import unittest
 import pygame
 import pygame.image
+import numpy as np
 
 from pgzero.screen import Screen
 from pgzero.loaders import set_root, images
@@ -11,28 +12,54 @@
 
 class ScreenTest(unittest.TestCase):
     @classmethod
-    def setUpClass(self):
+    def setUpClass(cls):
+        """Initialise the display and set loaders to target the current dir."""
+        pygame.init()
+        cls.surf = pygame.display.set_mode((200, 100))
         set_root(__file__)
 
+    @classmethod
+    def tearDownClass(cls):
+        """Shut down the display."""
+        pygame.display.quit()
+
     def setUp(self):
-        surf.fill((0, 0, 0))
-        self.screen = Screen(surf)
+        self.screen = Screen(self.surf)
+        self.screen.clear()
+
+    def assertImagesAlmostEqual(self, computed, expected):
+        """Check that 2 images are approximately equal."""
+        comp_surf = pygame.surfarray.array3d(computed)
+        exp_surf = pygame.surfarray.array3d(expected)
+
+        if np.allclose(comp_surf, exp_surf, atol=2):
+            return
 
-    def assertImagesEqual(self, a, b):
-        adata, bdata = (pygame.image.tostring(i, 'RGB') for i in (a, b))
+        name = sys._getframe(1).f_code.co_name
+        tmpdir = Path(__file__).parent / 'failed-image'
+        tmpdir.mkdir(exist_ok=True)
+        pygame.image.save(
+            computed,
+            str(tmpdir / '{}-computed.png'.format(name))
+        )
+        pygame.image.save(
+            expected,
+            str(tmpdir / '{}-expected.png'.format(name))
+        )
 
-        if adata != bdata:
-            raise AssertionError("Images differ")
+        raise AssertionError(
+            "Images differ; saved comparison images to {}".format(tmpdir)
+        )
 
     def test_blit_surf(self):
         """We can blit a surface to the screen."""
         self.screen.blit(images.alien, (0, 0))
-        self.assertImagesEqual(surf, images.expected_alien_blit)
+        self.assertImagesAlmostEqual(surf, images.expected_alien_blit)
 
     def test_blit_name(self):
         """screen.blit() accepts an image name instead of a Surface."""
         self.screen.blit('alien', (0, 0))
-        self.assertImagesEqual(surf, images.expected_alien_blit)
+        self.assertImagesAlmostEqual(surf, images.expected_alien_blit)
 
 
 if __name__ == '__main__':