File size: 5,646 Bytes
71e71e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | import logging
from datetime import datetime
from src.tg_compat.ext import CallbackContext, ContextTypes
from resources import phrases
from src.model.enums.ContextDataKey import ContextDataKey, ContextDataType
from src.model.error.CommonChatError import CommonChatException
from src.utils.download_utils import get_random_string
def get_context_data(
context: ContextTypes.DEFAULT_TYPE,
data_type: ContextDataType,
key: ContextDataKey,
inner_key: str = None,
tolerate_key_exception: bool = True,
) -> any:
"""
Get the bot context data
:param context: The context
:param data_type: The type
:param key: The key
:param inner_key: The inner key
:param tolerate_key_exception: If the key is not found, raise CommonChatException instead of
KeyError
:return: The data
"""
data = context.bot_data if data_type is ContextDataType.BOT else context.user_data
try:
if inner_key is not None:
return data[key]["value"][inner_key]["value"]
else:
return data[key]["value"]
except KeyError as e:
if tolerate_key_exception:
logging.debug(
f"Key {key} not found in context user data, this might be due to a system restart"
)
raise CommonChatException(phrases.RESTART_BOT, only_message=True)
else:
raise e
def set_context_data(
context: ContextTypes.DEFAULT_TYPE,
data_type: ContextDataType,
key: ContextDataKey,
value: any,
inner_key: str = None,
) -> None:
"""
Set the bot context data
:param context: The context
:param data_type: The type
:param key: The key
:param value: The value
:param inner_key: The inner key
:return: None
"""
data = context.bot_data if data_type is ContextDataType.BOT else context.user_data
if key not in data:
data[key] = {"value": {}}
if inner_key is not None:
data[key]["value"][inner_key] = {"value": value}
data[key]["value"][inner_key]["last_updated"] = datetime.now()
else:
data[key]["value"] = value
data[key]["last_updated"] = datetime.now()
def remove_context_data(
context: ContextTypes.DEFAULT_TYPE,
data_type: ContextDataType,
key: ContextDataKey,
inner_key: str = None,
) -> None:
"""
Remove the bot context data
:param context: The context
:param data_type: The type
:param key: The key
:param inner_key: The inner key
:return: None
"""
data = context.bot_data if data_type is ContextDataType.BOT else context.user_data
if key not in data:
return
if inner_key is not None:
data[key]["value"].pop(inner_key)
else:
data.pop(key)
def remove_user_context_data(
context: CallbackContext, key: ContextDataKey, inner_key: str = None
) -> None:
"""
Remove the user context data
:param context: The context
:param key: The key
:param inner_key: The inner key
:return: None
"""
remove_context_data(context, ContextDataType.USER, key, inner_key)
def get_bot_context_data(
context: CallbackContext,
key: ContextDataKey,
inner_key: str = None,
tolerate_key_exception: bool = True,
) -> any:
"""
Get the bot context data
:param context: The context
:param key: The key
:param inner_key: The inner key
:param tolerate_key_exception: If the key is not found, raise CommonChatException instead of
KeyError
:return: The data
"""
return get_context_data(context, ContextDataType.BOT, key, inner_key, tolerate_key_exception)
def set_bot_context_data(
context: CallbackContext, key: ContextDataKey, value: any, inner_key: str = None
) -> None:
"""
Set the bot context data
:param context: The context
:param key: The key
:param value: The value
:param inner_key: The inner key
:return: None
"""
set_context_data(context, ContextDataType.BOT, key, value, inner_key)
def get_user_context_data(
context: CallbackContext,
key: ContextDataKey,
inner_key: str = None,
tolerate_key_exception: bool = True,
) -> any:
"""
Get the user context data
:param context: The context
:param key: The key
:param inner_key: The inner key
:param tolerate_key_exception: If the key is not found, raise CommonChatException instead of
KeyError
:return: The data
"""
return get_context_data(context, ContextDataType.USER, key, inner_key, tolerate_key_exception)
def set_user_context_data(
context: CallbackContext, key: ContextDataKey, value: any, inner_key: str = None
) -> None:
"""
Set the user context data
:param context: The context
:param key: The key
:param value: The value
:param inner_key: The inner key
:return: None
"""
set_context_data(context, ContextDataType.USER, key, value, inner_key)
def get_random_user_context_inner_query_key(context: CallbackContext) -> str:
"""
Get a random inner key for the user context data
:param context: The context
:return: The inner key
"""
for i in range(100):
inner_key = get_random_string()
try:
get_user_context_data(
context,
ContextDataKey.INLINE_QUERY,
inner_key=inner_key,
tolerate_key_exception=False,
)
continue
except KeyError:
return inner_key
logging.error("Could not find a random inner key for the user context data")
raise CommonChatException("Failed to generate an inner key")
|