Add support for State type in scenes methods (#1685)

* Add support for `State` type to scenes methods like `goto`, `enter`, `get`

Calm down MyPy

* Add 1685.feature.rst

CHANGES/1685.feature.rst

@@ -0,0 +1 @@
+Add support for `State` type in scenes methods like `goto`, `enter`, `get`

aiogram/fsm/scene.py

@@ -172,7 +172,7 @@ class ActionContainer:
         name: str,
         filters: Tuple[CallbackType, ...],
         action: SceneAction,
-        target: Optional[Union[Type[Scene], str]] = None,
+        target: Optional[Union[Type[Scene], State, str]] = None,
     ) -> None:
         self.name = name
         self.filters = filters
@@ -537,14 +537,14 @@ class SceneWizard:
         assert self.scene_config.state is not None, "Scene state is not specified"
         await self.goto(self.scene_config.state, **kwargs)
 
-    async def goto(self, scene: Union[Type[Scene], str], **kwargs: Any) -> None:
+    async def goto(self, scene: Union[Type[Scene], State, str], **kwargs: Any) -> None:
         """
         The `goto` method transitions to a new scene.
         It first calls the `leave` method to perform any necessary cleanup
         in the current scene, then calls the `enter` event to enter the specified scene.
 
         :param scene: The scene to transition to. Can be either a `Scene` instance
-            or a string representing the scene.
+            `State` instance or a string representing the scene.
         :param kwargs: Additional keyword arguments to pass to the `enter`
             method of the scene manager.
         :return: None
@@ -665,7 +665,7 @@ class ScenesManager:
 
         self.history = HistoryManager(self.state)
 
-    async def _get_scene(self, scene_type: Optional[Union[Type[Scene], str]]) -> Scene:
+    async def _get_scene(self, scene_type: Optional[Union[Type[Scene], State, str]]) -> Scene:
         scene_type = self.registry.get(scene_type)
         return scene_type(
             wizard=SceneWizard(
@@ -687,14 +687,14 @@ class ScenesManager:
 
     async def enter(
         self,
-        scene_type: Optional[Union[Type[Scene], str]],
+        scene_type: Optional[Union[Type[Scene], State, str]],
         _check_active: bool = True,
         **kwargs: Any,
     ) -> None:
         """
         Enters the specified scene.
 
-        :param scene_type: Optional Type[Scene] or str representing the scene type to enter.
+        :param scene_type: Optional Type[Scene], State or str representing the scene type to enter.
         :param _check_active: Optional bool indicating whether to check if
             there is an active scene to exit before entering the new scene. Defaults to True.
         :param kwargs: Additional keyword arguments to pass to the scene's wizard.enter() method.
@@ -833,18 +833,19 @@ class SceneRegistry:
         """
         self.add(*scenes, router=self.router)
 
-    def get(self, scene: Optional[Union[Type[Scene], str]]) -> Type[Scene]:
+    def get(self, scene: Optional[Union[Type[Scene], State, str]]) -> Type[Scene]:
         """
         This method returns the registered Scene object for the specified scene.
-        The scene parameter can be either a Scene object or a string representing
+        The scene parameter can be either a Scene object, State object or a string representing
         the name of the scene. If a Scene object is provided, the state attribute
         of the SceneConfig object associated with the Scene object will be used as the scene name.
-        If None or an invalid type is provided, a SceneException will be raised.
+        If a State object is provided, the state attribute of the State object will be used as the
+        scene name. If None or an invalid type is provided, a SceneException will be raised.
 
         If the specified scene is not registered in the SceneRegistry object,
         a SceneException will be raised.
 
-        :param scene: A Scene object or a string representing the name of the scene.
+        :param scene: A Scene object, State object or a string representing the name of the scene.
         :return: The registered Scene object corresponding to the given scene parameter.
 
         """
@@ -853,7 +854,7 @@ class SceneRegistry:
         if isinstance(scene, State):
             scene = scene.state
         if scene is not None and not isinstance(scene, str):
-            raise SceneException("Scene must be a subclass of Scene or a string")
+            raise SceneException("Scene must be a subclass of Scene, State or a string")
 
         try:
             return self._scenes[scene]
@@ -864,7 +865,7 @@ class SceneRegistry:
 @dataclass
 class After:
     action: SceneAction
-    scene: Optional[Union[Type[Scene], str]] = None
+    scene: Optional[Union[Type[Scene], State, str]] = None
 
     @classmethod
     def exit(cls) -> After:
@@ -875,7 +876,7 @@ class After:
         return cls(action=SceneAction.back)
 
     @classmethod
-    def goto(cls, scene: Optional[Union[Type[Scene], str]]) -> After:
+    def goto(cls, scene: Optional[Union[Type[Scene], State, str]]) -> After:
         return cls(action=SceneAction.enter, scene=scene)
 
 

tests/test_fsm/test_scene.py

@@ -1490,7 +1490,9 @@ class TestSceneRegistry:
         registry = SceneRegistry(router, register_on_add)
         registry.add(MyScene)
 
-        with pytest.raises(SceneException, match="Scene must be a subclass of Scene or a string"):
+        with pytest.raises(
+            SceneException, match="Scene must be a subclass of Scene, State or a string"
+        ):
             registry.get(MyScene)
 
     def test_scene_registry_get_scene_not_registered(self):