@@ -1301,3 +1301,6 @@ msgstr "" msgid "You cannot change the password of an unconfirmed e-mail user." msgstr "Вы не можете изменить пароль неподтвержденного по e-mail пользователя." + +msgid "Unknown file format" +msgstr "Неизвестный формат файла" @@ -36,7 +36,7 @@ class MessageSerializer(serializers.ModelSerializer): ] def validate_file(self, file: UploadedFile | None) -> UploadedFile: - max_mb_size = 8 + max_mb_size = 5 if file and file.size > (max_mb_size << 10 << 10): raise ValidationError( _('The file size cannot exceed %(max_mb_size)d MB') % {'max_mb_size': max_mb_size} @@ -2,6 +2,7 @@ import base64 import logging import subprocess import time +import zipfile from datetime import timedelta from decimal import Decimal from io import BytesIO, StringIO @@ -10,6 +11,7 @@ from typing import Any, Callable, Literal from uuid import UUID import docx2txt +import openpyxl import filetype from django.core.cache import cache from django.core.files.base import ContentFile @@ -23,7 +25,7 @@ from messages.models import Message from ml_model.exceptions import ( InferenceDisabled, PaymentRuleNotImplemented, - ScraperDoesNotExists, + ScraperDoesNotExists, UnknownFileException, ) from ml_model.models import ( Deployment, @@ -118,6 +120,19 @@ class InferenceService: raw_file.seek(0) file_buf = BytesIO(raw_file.read()) + if file_extension == 'zip': + file_extension = None + signatures = { + 'xlsx': 'xl/workbook.xml' + } + with zipfile.ZipFile(file_buf, 'r') as zip_file: + namelist = zip_file.namelist() + for format_name, required_file in signatures.items(): + if required_file in namelist: + file_extension = format_name + break + if not file_extension: + raise UnknownFileException if file_extension in ('png', 'jpg', 'jpeg'): file = BytesIO() normalized_image = ImageModule.open(file_buf) @@ -128,7 +143,7 @@ class InferenceService: reader = PdfReader(file_buf) for page in reader.pages: file.write(page.extract_text()) - elif file_extension in ('doc', 'docx', 'zip'): + elif file_extension in ('doc', 'docx'): extractors: dict[Literal['doc', 'docx'], Callable[[], str]] = { 'doc': lambda: subprocess.Popen( ['antiword', '-w', '0', '-'], @@ -139,13 +154,20 @@ class InferenceService: .communicate(file_buf.getvalue()) .decode(), 'docx': lambda: docx2txt.process(file_buf), - 'zip': lambda: docx2txt.process(file_buf), } file = StringIO() file.write('Remember this Document included in request:') file.write('[DOCUMENT-START]\n') file.write(extractors[file_extension]()) file.write('\n[DOCUMENT-END]') + elif file_extension in ('xlsx',): + file = StringIO() + xlsx_file = openpyxl.load_workbook(file_buf) + for sheet_name in xlsx_file.sheetnames: + sheet = xlsx_file[sheet_name] + for row in sheet.iter_rows(values_only=True): + file.write(f'Данные ряда: {row}') + except Exception: file = None @@ -29,3 +29,8 @@ class PaymentRuleNotImplemented(Exception): class ScraperDoesNotExists(Exception): def __str__(self): return _('Scraper does not exists') + + +class UnknownFileException(Exception): + def __str__(self): + return _('Unknown file format') \ No newline at end of file