classes/session_document_store.py

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_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