paulmataruso/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2025-11-13 02:17:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

embedded bots: Support multi-entry state modification.

Browse Source
This commit is contained in:
derAnfaenger
2017-11-20 14:54:41 +01:00
committed by Tim Abbott
parent cbadd3d430
commit d2af8d4cbd
2 changed files with 22 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
zerver/lib/bot_lib.py
View File

@@ -49,9 +49,10 @@ class StateHandler:
def put(self, key: Text, value: Text) -> None: def put(self, key: Text, value: Text) -> None:
set_bot_state(self.user_profile, key, self.marshal(value)) set_bot_state(self.user_profile, key, self.marshal(value))
set_bot_state(self.user_profile, [(key, self.marshal(value))])
def remove(self, key: Text) -> None: def remove(self, key: Text) -> None:
remove_bot_state(self.user_profile, key) remove_bot_state(self.user_profile, [key])
def contains(self, key: Text) -> bool: def contains(self, key: Text) -> bool:
return is_key_in_bot_state(self.user_profile, key) return is_key_in_bot_state(self.user_profile, key)

42
zerver/lib/bot_storage.py
View File

@@ -4,7 +4,7 @@ from django.db.models.query import F
from django.db.models.functions import Length from django.db.models.functions import Length
from zerver.models import BotUserStateData, UserProfile, Length from zerver.models import BotUserStateData, UserProfile, Length
from typing import Text, Optional from typing import Text, Optional, List, Tuple
class StateError(Exception): class StateError(Exception):
pass pass
@@ -28,33 +28,31 @@ def get_bot_state_size(bot_profile, key=None):
except BotUserStateData.DoesNotExist: except BotUserStateData.DoesNotExist:
return 0 return 0
def set_bot_state(bot_profile, key, value): def set_bot_state(bot_profile, entries):
# type: (UserProfile, Text, Text) -> None # type: (UserProfile, List[Tuple[str, str]]) -> None
state_size_limit = settings.USER_STATE_SIZE_LIMIT state_size_limit = settings.USER_STATE_SIZE_LIMIT
old_entry_size = get_bot_state_size(bot_profile, key) state_size_difference = 0
new_entry_size = len(key) + len(value) for key, value in entries:
old_state_size = get_bot_state_size(bot_profile) if type(key) is not str:
new_state_size = old_state_size + (new_entry_size - old_entry_size) raise StateError("Key type is {}, but should be str.".format(type(key)))
if type(value) is not str:
raise StateError("Value type is {}, but should be str.".format(type(value)))
state_size_difference += (len(key) + len(value)) - get_bot_state_size(bot_profile, key)
new_state_size = get_bot_state_size(bot_profile) + state_size_difference
if new_state_size > state_size_limit: if new_state_size > state_size_limit:
raise StateError("Request exceeds storage limit by {} characters. The limit is {} characters." raise StateError("Request exceeds storage limit by {} characters. The limit is {} characters."
.format(new_state_size - state_size_limit, state_size_limit)) .format(new_state_size - state_size_limit, state_size_limit))
elif type(key) is not str:
raise StateError("Key type is {}, but should be str.".format(type(key)))
elif type(value) is not str:
raise StateError("Value type is {}, but should be str.".format(type(value)))
else: else:
obj, created = BotUserStateData.objects.get_or_create(bot_profile=bot_profile, key=key, for key, value in entries:
defaults={'value': value}) BotUserStateData.objects.update_or_create(bot_profile=bot_profile, key=key,
if not created: defaults={'value': value})
obj.value = value
obj.save()
def remove_bot_state(bot_profile, key): def remove_bot_state(bot_profile, keys):
# type: (UserProfile, Text) -> None # type: (UserProfile, List[Text]) -> None
try: queryset = BotUserStateData.objects.filter(bot_profile=bot_profile, key__in=keys)
BotUserStateData.objects.get(bot_profile=bot_profile, key=key).delete() if len(queryset) < len(keys):
except BotUserStateData.DoesNotExist: raise StateError("Key does not exist.")
raise StateError("Key does not exist.".format(key)) queryset.delete()
def is_key_in_bot_state(bot_profile, key): def is_key_in_bot_state(bot_profile, key):
# type: (UserProfile, Text) -> bool # type: (UserProfile, Text) -> bool