Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 2,678 Bytes
cbfe36d | 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 | import sys
from typing import Dict, List, Tuple
from constants import MAX_FILE_SIZES_PER_SESSION
class SessionDocumentStore:
def __init__(self) -> None:
# Stores, for each session, the files' content and name
# session_id -> {file_name -> (file_text, size_in_bytes)}
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_documents(self, session_id: str) -> Dict[str, str] | None:
"""Return {file_name: file_text} for the session, or None if it has none."""
document_map = self.session_document_map.get(session_id)
if not document_map:
return None
return {name: text for name, (text, _size) in document_map.items()}
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
|