Minor changes

1 files changed, 12 insertions, 7 deletions

diff --git a/snake.py b/snake.py
index ca21129..a37fdf6 100755
--- a/snake.py
+++ b/snake.py
@@ -129,7 +129,7 @@ class Snake:
         self.draw_infos()
         pygame.display.flip()
 
-    def play(self,direction):
+    def play(self,direction,handle_quit=True):
         """
         Play using wall clock directions (12=up, 3=right, 6=down and 9=left)
         """
@@ -151,10 +151,11 @@ class Snake:
         self.draw()
         self.clock.tick(self.fps)
         # Ensure we not quit
-        for event in pygame.event.get():
-                if event.type == pygame.QUIT: 
-                    pygame.quit()
-                    sys.exit()
+        if handle_quit:
+            for event in pygame.event.get():
+                    if event.type == pygame.QUIT: 
+                        pygame.quit()
+                        sys.exit()
         return(code)
     
     def play2(self,direction):
@@ -190,7 +191,10 @@ class Snake:
         while True:
             # Check inputs
             for event in pygame.event.get():
-                if event.type == pygame.KEYDOWN:
+                if event.type == pygame.QUIT: 
+                    pygame.quit()
+                    sys.exit()
+                elif event.type == pygame.KEYDOWN:
                     if event.key == pygame.K_LEFT and self.direction != 3:
                         self.direction=9
                         break
@@ -204,5 +208,6 @@ class Snake:
                         self.direction=6
                         break
             
-            if self.play(self.direction) <0:
+            if self.play(self.direction,handle_quit=False) <0:
                 break
+            
\ No newline at end of file