fix: add image handler exception (#27)

This commit is contained in:
Taeho 2023-04-08 03:51:09 +09:00 committed by GitHub
parent 3b5e7f4872
commit 51e351904a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -86,6 +86,17 @@ class FileHandler:
shutil.copy(src, dst) shutil.copy(src, dst)
else: else:
local_filename = self.download(url) local_filename = self.download(url)
return self.handlers[FileType.from_url(url)].handle(local_filename)
try:
handler = self.handlers[FileType.from_url(url)]
except KeyError:
if FileType.from_url(url) == FileType.IMAGE:
raise Exception(
f"No handler for {FileType.from_url(url)}. "
f"Please set USE_GPU to True in env/settings.py"
)
else:
raise Exception(f"No handler for {FileType.from_url(url)}")
handler.handle(local_filename)
except Exception as e: except Exception as e:
raise e raise e