From 90b3a990396dc65ecc1029997a349c5555ea382c Mon Sep 17 00:00:00 2001
From: Andrey Tikhonov <andrey.tikhonov@quantori.com>
Date: Sun, 5 Sep 2021 23:55:38 +0300
Subject: [PATCH] iter states in states group (#666)

* iter states in states group

* fix type hint

* remove empty line

* add changes for doc
---
 CHANGES/666.feature                          | 2 ++
 aiogram/dispatcher/fsm/state.py              | 5 ++++-
 tests/test_dispatcher/test_fsm/test_state.py | 7 +++++++
 3 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 CHANGES/666.feature

diff --git a/CHANGES/666.feature b/CHANGES/666.feature
new file mode 100644
index 00000000..35a6572e
--- /dev/null
+++ b/CHANGES/666.feature
@@ -0,0 +1,2 @@
+Ability to iterate over all states in StatesGroup.
+Aiogram already had in check for states group so this is relative feature.
diff --git a/aiogram/dispatcher/fsm/state.py b/aiogram/dispatcher/fsm/state.py
index ced9779a..a034e003 100644
--- a/aiogram/dispatcher/fsm/state.py
+++ b/aiogram/dispatcher/fsm/state.py
@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Optional, Tuple, Type, no_type_check
+from typing import Any, Iterator, Optional, Tuple, Type, no_type_check
 
 from ...types import TelegramObject
 
@@ -118,6 +118,9 @@ class StatesGroupMeta(type):
     def __str__(self) -> str:
         return f"<StatesGroup '{self.__full_group_name__}'>"
 
+    def __iter__(self) -> Iterator[State]:
+        return iter(self.__all_states__)
+
 
 class StatesGroup(metaclass=StatesGroupMeta):
     @classmethod
diff --git a/tests/test_dispatcher/test_fsm/test_state.py b/tests/test_dispatcher/test_fsm/test_state.py
index 04c8e448..3d41e1fa 100644
--- a/tests/test_dispatcher/test_fsm/test_state.py
+++ b/tests/test_dispatcher/test_fsm/test_state.py
@@ -150,6 +150,13 @@ class TestStatesGroup:
 
         assert MyGroup.MyNestedGroup.get_root() is MyGroup
 
+    def test_iterable(self):
+        class Group(StatesGroup):
+            x = State()
+            y = State()
+
+        assert set(Group) == {Group.x, Group.y}
+
     def test_empty_filter(self):
         class MyGroup(StatesGroup):
             pass