Adding use_bot_state

Author: axelsrz

libraries/botbuilder-core/botbuilder/core/adapter_extensions.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 from botbuilder.core import (
     BotAdapter,
+    BotState,
     Storage,
     RegisterClassMiddleware,
     UserState,
@@ -23,6 +24,39 @@ def use_storage(adapter: BotAdapter, storage: Storage) -> BotAdapter:
         """
         return adapter.use(RegisterClassMiddleware(storage))
 
+    @staticmethod
+    def use_bot_state(
+        bot_adapter: BotAdapter, *bot_states: BotState, auto: bool = True
+    ) -> BotAdapter:
+        """
+        Registers bot state object into the TurnContext. The botstate will be available via the turn context.
+
+        :param bot_adapter: The BotAdapter on which to register the state objects.
+        :param bot_states: One or more BotState objects to register.
+        :return: The updated adapter.
+        """
+        if not bot_states:
+            raise TypeError("At least one BotAdapter is required")
+
+        for bot_state in bot_states:
+            bot_adapter.use(
+                RegisterClassMiddleware(
+                    bot_state, AdapterExtensions.fullname(bot_state)
+                )
+            )
+
+        if auto:
+            bot_adapter.use(AutoSaveStateMiddleware(bot_states))
+
+        return bot_adapter
+
+    @staticmethod
+    def fullname(obj):
+        module = obj.__class__.__module__
+        if module is None or module == str.__class__.__module__:
+            return obj.__class__.__name__  # Avoid reporting __builtin__
+        return module + "." + obj.__class__.__name__
+
     @staticmethod
     def use_state(
         adapter: BotAdapter,

libraries/botbuilder-core/botbuilder/core/register_class_middleware.py

@@ -10,16 +10,18 @@ class RegisterClassMiddleware(Middleware):
     Middleware for adding an object to or registering a service with the current turn context.
     """
 
-    def __init__(self, service):
+    def __init__(self, service, key: str = None):
         self.service = service
+        self._key = key
 
     async def on_turn(
         self, context: TurnContext, logic: Callable[[TurnContext], Awaitable]
     ):
         # C# has TurnStateCollection with has overrides for adding items
         # to TurnState.  Python does not.  In C#'s case, there is an 'Add'
         # to handle adding object, and that uses the fully qualified class name.
-        context.turn_state[self.fullname(self.service)] = self.service
+        key = self._key or self.fullname(self.service)
+        context.turn_state[key] = self.service
         await logic()
 
     @staticmethod

libraries/botbuilder-dialogs/tests/test_dialogextensions.py

@@ -221,7 +221,7 @@ async def capture_eoc(
             logic, TestAdapter.create_conversation_reference(conversation_id)
         )
         AdapterExtensions.use_storage(adapter, storage)
-        AdapterExtensions.use_state(adapter, user_state, convo_state)
+        AdapterExtensions.use_bot_state(adapter, user_state, convo_state)
         adapter.use(TranscriptLoggerMiddleware(ConsoleTranscriptLogger()))
 
         return TestFlow(None, adapter)