| import sys
|
|
|
| from typing import Dict, List, Tuple
|
| from constants import MAX_FILE_SIZES_PER_SESSION
|
|
|
|
|
| class SessionDocumentStore:
|
| def __init__(self) -> None:
|
|
|
|
|
| self.session_document_map: Dict[str, Dict[str, Tuple[str, int]]] = dict()
|
|
|
| def create_document(self, session_id: str, file_text: str, file_name: str):
|
| text_size = sys.getsizeof(file_text)
|
| if session_id not in self.session_document_map:
|
| if text_size > MAX_FILE_SIZES_PER_SESSION:
|
| return False
|
| self.session_document_map[session_id] = dict()
|
| self.session_document_map[session_id][file_name] = (file_text, text_size)
|
| return True
|
|
|
| current_total_file_size = sum(
|
| file_text_size[1]
|
| for file_text_size in self.session_document_map[session_id].values()
|
| )
|
|
|
| if current_total_file_size + text_size > MAX_FILE_SIZES_PER_SESSION:
|
| return False
|
|
|
| self.session_document_map[session_id][file_name] = (file_text, text_size)
|
| return True
|
|
|
| def get_document_contents(self, session_id: str) -> List[str] | None:
|
| document_map = self.session_document_map.get(session_id)
|
| if document_map is None:
|
| return None
|
|
|
| document_contents = [
|
| file_text_size[0] for file_text_size in document_map.values()
|
| ]
|
| if len(document_contents) == 0:
|
| return None
|
|
|
| return document_contents
|
|
|
| def delete_document(self, session_id: str, file_name: str) -> bool:
|
| """Deletes a document with the given name. If the session no longer has documents
|
| after the deletion, the session is also deleted and the function returns True."""
|
| document_map = self.session_document_map.get(session_id)
|
| if document_map is None:
|
| return False
|
|
|
| if file_name in document_map:
|
| del document_map[file_name]
|
| if len(document_map) == 0:
|
| self.delete_session_documents(session_id)
|
| return True
|
|
|
| return False
|
|
|
| def delete_session_documents(self, session_id: str) -> bool:
|
| return self.session_document_map.pop(session_id, None) is not None
|
|
|
