Team Ai
Apppublic

kumarAnurag/ocr_image_file_processing

sourceHugging Faceupdated 2y agoView on Hugging Face
1likes
test_loader.py56 linesDownload Raw Back to root
1import os2 3# Mock loader classes for demonstration purposes4class PyPDFLoader:5    def __init__(self, file_path):6        self.file_path = file_path7        print(f'PDF loader initialized with {file_path}')8 9class TextLoader:10    def __init__(self, file_path):11        self.file_path = file_path12        print(f'Text loader initialized with {file_path}')13 14class CSVLoader:15    def __init__(self, file_path):16        self.file_path = file_path17        print(f'CSV loader initialized with {file_path}')18 19# Function to determine the loader based on file extension20def get_loader_by_file_extension(temp_file):21    if not isinstance(temp_file, str):22        raise TypeError("Expected file path as a string.")23    24    file_split = os.path.splitext(temp_file)25    file_extension = file_split[1]  # Extract the extension26    print('file_extension - ', file_extension)27 28    # Initialize loader based on file extension29    if file_extension == '.pdf':30        loader = PyPDFLoader(temp_file)31        print('Loader Created for PDF file')32    elif file_extension == '.txt':33        loader = TextLoader(temp_file)34    elif file_extension == '.csv':35        loader = CSVLoader(temp_file)36    else:37        raise ValueError(f"Unsupported file type: {file_extension}")38 39    return loader40 41# Test the function with different file types42if __name__ == "__main__":43    test_files = [44        "document.pdf",45        "notes.txt",46        "data.csv",47        "image.jpg"  # This should raise an error48    ]49 50    for test_file in test_files:51        print(f"\nTesting with file: {test_file}")52        try:53            loader = get_loader_by_file_extension(test_file)54        except Exception as e:55            print(f'Error: {e}')56