Repository: MyreMylar/pygame_gui
Pull request #546 (2024-03-16 12:08:13)

Add iteration support to IUIContainerInterface


Patch:
diff --git a/pygame_gui/core/interfaces/container_interface.py b/pygame_gui/core/interfaces/container_interface.py
index 497813b..0537939 100644
--- a/pygame_gui/core/interfaces/container_interface.py
+++ b/pygame_gui/core/interfaces/container_interface.py
@@ -1,3 +1,4 @@
+import typing
 from abc import ABCMeta, abstractmethod
 from typing import Tuple, Union
 
@@ -186,6 +187,13 @@ class IUIContainerInterface(IUIElementInterface, metaclass=ABCMeta):
         :param element: the element to check.
         """
 
+    @abstractmethod
+    def __iter__(self) -> typing.Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container-like interface.
+        This method allows iterating over the children elements within the container.
+        """
+
 
 class IContainerLikeInterface(metaclass=ABCMeta):
     """
@@ -217,3 +225,10 @@ class IContainerLikeInterface(metaclass=ABCMeta):
         process events. Should also hide all the children elements.
         If the container was hidden before - ignore.
         """
+        
+    @abstractmethod
+    def __iter__(self) -> typing.Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container-like interface. 
+        This method allows iterating over the children elements within the container.
+        """
diff --git a/pygame_gui/core/ui_container.py b/pygame_gui/core/ui_container.py
index 03047f7..c6f33b2 100644
--- a/pygame_gui/core/ui_container.py
+++ b/pygame_gui/core/ui_container.py
@@ -1,4 +1,4 @@
-from typing import List, Union, Tuple, Dict
+from typing import List, Union, Tuple, Dict, Iterator
 
 import pygame
 
@@ -395,3 +395,9 @@ class UIContainer(UIElement, IUIContainerInterface, IContainerLikeInterface):
             if target in element.get_anchor_targets():
                 element.update_containing_rect_position()
                 self.on_contained_elements_changed(element)
+                
+    def __iter__(self) -> Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container.
+        """
+        return iter(self.elements)
diff --git a/pygame_gui/elements/ui_panel.py b/pygame_gui/elements/ui_panel.py
index 3830cf0..5e400ae 100644
--- a/pygame_gui/elements/ui_panel.py
+++ b/pygame_gui/elements/ui_panel.py
@@ -1,9 +1,9 @@
-from typing import Union, Dict, Tuple, Optional
+from typing import Union, Dict, Tuple, Optional, Iterator
 
 import pygame
 
 from pygame_gui.core import ObjectID
-from pygame_gui.core.interfaces import IUIManagerInterface
+from pygame_gui.core.interfaces import IUIManagerInterface, IUIElementInterface
 from pygame_gui.core.interfaces import IContainerLikeInterface, IUIContainerInterface
 
 from pygame_gui.core import UIElement, UIContainer
@@ -308,3 +308,9 @@ class UIPanel(UIElement, IContainerLikeInterface):
         if self.panel_container is not None:
             self.panel_container.hide()
         super().hide()
+    
+    def __iter__(self) -> Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container.
+        """
+        return iter(self.get_container())
diff --git a/pygame_gui/elements/ui_scrolling_container.py b/pygame_gui/elements/ui_scrolling_container.py
index 80a91c3..d0794c1 100644
--- a/pygame_gui/elements/ui_scrolling_container.py
+++ b/pygame_gui/elements/ui_scrolling_container.py
@@ -1,10 +1,10 @@
-from typing import Union, Dict, Tuple, Optional, List
+from typing import Union, Dict, Tuple, Optional, List, Iterator
 
 import pygame
 
 from pygame_gui.core import ObjectID
 from pygame_gui.core.interfaces import IContainerLikeInterface, IUIContainerInterface
-from pygame_gui.core.interfaces import IUIManagerInterface, Coordinate
+from pygame_gui.core.interfaces import IUIManagerInterface, Coordinate, IUIElementInterface
 from pygame_gui.core import UIElement, UIContainer
 
 from pygame_gui.elements.ui_vertical_scroll_bar import UIVerticalScrollBar
@@ -469,3 +469,9 @@ class UIScrollingContainer(UIElement, IContainerLikeInterface):
         if self._root_container is not None:
             self._root_container.hide()
         super().hide()
+        
+    def __iter__(self) -> Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container.
+        """
+        return iter(self.get_container())
diff --git a/pygame_gui/elements/ui_window.py b/pygame_gui/elements/ui_window.py
index 89eb7d4..463289c 100644
--- a/pygame_gui/elements/ui_window.py
+++ b/pygame_gui/elements/ui_window.py
@@ -1,4 +1,4 @@
-from typing import Union, Tuple, Optional, List
+from typing import Union, Tuple, Optional, List, Iterator
 
 import pygame
 
@@ -7,7 +7,7 @@ from pygame_gui._constants import OldType
 
 from pygame_gui.core import ObjectID
 from pygame_gui.core.interfaces import IContainerLikeInterface, IUIContainerInterface
-from pygame_gui.core.interfaces import IWindowInterface, IUIManagerInterface
+from pygame_gui.core.interfaces import IWindowInterface, IUIManagerInterface, IUIElementInterface
 from pygame_gui.core import UIElement, UIContainer
 from pygame_gui.core.drawable_shapes import RectDrawableShape, RoundedRectangleShape
 
@@ -753,3 +753,9 @@ class UIWindow(UIElement, IContainerLikeInterface, IWindowInterface):
                              abs_mouse_pos[1] - window_contents_top_left[1])
 
         return rel_mouse_pos
+    
+    def __iter__(self) -> Iterator[IUIElementInterface]:
+        """
+        Iterates over the elements within the container.
+        """
+        return iter(self.get_container())


Test patch:
diff --git a/tests/test_core/test_ui_container.py b/tests/test_core/test_ui_container.py
index cce27eb..7574d2a 100644
--- a/tests/test_core/test_ui_container.py
+++ b/tests/test_core/test_ui_container.py
@@ -264,6 +264,20 @@ class TestUIContainer:
         assert button.get_relative_rect() == pygame.Rect(50, 100, 50, 50)
         assert container.get_relative_rect() == pygame.Rect(100, 50, 200, 250)
 
+    def test_iteration(self, _init_pygame, default_ui_manager: IUIManagerInterface,
+                       _display_surface_return_none):
+        container = UIContainer(pygame.Rect(100, 100, 200, 200), manager=default_ui_manager)
+        button_1 = UIButton(relative_rect=pygame.Rect(50, 50, 50, 50), text="1",
+                            manager=default_ui_manager, container=container)
+        button_2 = UIButton(relative_rect=pygame.Rect(150, 50, 50, 50), text="2",
+                            manager=default_ui_manager, container=container)
+
+        count = 0
+        for button in container:
+            button.get_relative_rect()
+            count += 1
+        assert count == 2
+
 
 if __name__ == '__main__':
     pytest.console_main()
diff --git a/tests/test_elements/test_ui_panel.py b/tests/test_elements/test_ui_panel.py
index 2fdbc37..b3d9a6b 100644
--- a/tests/test_elements/test_ui_panel.py
+++ b/tests/test_elements/test_ui_panel.py
@@ -600,6 +600,20 @@ class TestUIPanel:
         manager.draw_ui(surface)
         assert compare_surfaces(empty_surface, surface)
 
+    def test_iteration(self, _init_pygame, default_ui_manager: IUIManagerInterface,
+                       _display_surface_return_none):
+        panel = UIPanel(pygame.Rect(100, 100, 200, 200), manager=default_ui_manager)
+        button_1 = UIButton(relative_rect=pygame.Rect(50, 50, 50, 50), text="1",
+                            manager=default_ui_manager, container=panel)
+        button_2 = UIButton(relative_rect=pygame.Rect(150, 50, 50, 50), text="2",
+                            manager=default_ui_manager, container=panel)
+
+        count = 0
+        for button in panel:
+            button.get_relative_rect()
+            count += 1
+        assert count == 2
+
 
 if __name__ == '__main__':
     pytest.console_main()
diff --git a/tests/test_elements/test_ui_scrolling_container.py b/tests/test_elements/test_ui_scrolling_container.py
index b68e11a..0fdedd6 100644
--- a/tests/test_elements/test_ui_scrolling_container.py
+++ b/tests/test_elements/test_ui_scrolling_container.py
@@ -389,6 +389,20 @@ class TestUIScrollingContainer:
         manager.draw_ui(surface)
         assert compare_surfaces(empty_surface, surface)
 
+    def test_iteration(self, _init_pygame, default_ui_manager: IUIManagerInterface,
+                       _display_surface_return_none):
+        container = UIScrollingContainer(pygame.Rect(100, 100, 200, 200), manager=default_ui_manager)
+        button_1 = UIButton(relative_rect=pygame.Rect(50, 50, 50, 50), text="1",
+                            manager=default_ui_manager, container=container)
+        button_2 = UIButton(relative_rect=pygame.Rect(150, 50, 50, 50), text="2",
+                            manager=default_ui_manager, container=container)
+
+        count = 0
+        for button in container:
+            button.get_relative_rect()
+            count += 1
+        assert count == 2
+
 
 if __name__ == '__main__':
     pytest.console_main()
diff --git a/tests/test_elements/test_ui_window.py b/tests/test_elements/test_ui_window.py
index 6ec3184..b52849b 100644
--- a/tests/test_elements/test_ui_window.py
+++ b/tests/test_elements/test_ui_window.py
@@ -667,6 +667,20 @@ class TestUIWindow:
 
         assert window.rect.width == 100 and window.rect.height == 100
 
+    def test_iteration(self, _init_pygame, default_ui_manager: IUIManagerInterface,
+                       _display_surface_return_none):
+        window = UIWindow(pygame.Rect(100, 100, 200, 200), manager=default_ui_manager)
+        button_1 = UIButton(relative_rect=pygame.Rect(50, 50, 50, 50), text="1",
+                            manager=default_ui_manager, container=window)
+        button_2 = UIButton(relative_rect=pygame.Rect(150, 50, 50, 50), text="2",
+                            manager=default_ui_manager, container=window)
+
+        count = 0
+        for button in window:
+            button.get_relative_rect()
+            count += 1
+        assert count == 2
+
 
 if __name__ == '__main__':
     pytest.console_main()