classes/session_document_store.py

import sys

from typing import Dict, List, Tuple
from constants import MAX_FILE_SIZE, MAX_FILE_SIZES_PER_SESSION
from exceptions import FileDocumentStoreError, FileDocumentStoreException


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_SIZE:
                raise FileDocumentStoreException(
                    FileDocumentStoreError.EXCEED_MAX_FILE_SIZE
                )
            self.session_document_map[session_id] = dict()
            self.session_document_map[session_id][file_name] = (file_text, text_size)
            return

        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:
            raise FileDocumentStoreException(
                FileDocumentStoreError.EXCEED_MAX_FILE_SIZE
            )

        if file_name in self.session_document_map[session_id]:
            raise FileDocumentStoreException(FileDocumentStoreError.FILE_ALREADY_EXISTS)

        self.session_document_map[session_id][file_name] = (file_text, text_size)

    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