File size: 682 Bytes
71e7dd8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import os
from typing import Type

from buster.documents.base import DocumentsManager
from buster.documents.pickle import DocumentsPickle
from buster.documents.sqlite import DocumentsDB

PICKLE_EXTENSIONS = [".gz", ".bz2", ".zip", ".xz", ".zst", ".tar", ".tar.gz", ".tar.xz", ".tar.bz2"]


def get_file_extension(filepath: str) -> str:
    return os.path.splitext(filepath)[1]


def get_documents_manager_from_extension(filepath: str) -> Type[DocumentsManager]:
    ext = get_file_extension(filepath)

    if ext in PICKLE_EXTENSIONS:
        return DocumentsPickle
    elif ext == ".db":
        return DocumentsDB
    else:
        raise ValueError(f"Unsupported format: {ext}.")