fixed single file handling, added keepfilename

This commit is contained in:
Jack Neil 2021-04-26 13:45:56 -04:00
parent c9011d066b
commit 3aede5d9b7
2 changed files with 12 additions and 6 deletions

View File

@ -22,6 +22,7 @@ from ffmpy import FFmpeg
import math
import tempfile
from pandas.api.types import is_bool_dtype, is_numeric_dtype, is_string_dtype
from pathlib import Path
class InputComponent(Component):
@ -889,16 +890,18 @@ class File(InputComponent):
Input type: Union[file-object, bytes, List[Union[file-object, bytes]]]
"""
def __init__(self, file_count="single", type="file", label=None):
def __init__(self, file_count="single", type="file", label=None, keepfilename=True):
'''
Parameters:
file_count (str): if single, allows user to upload one file. If "multiple", user uploads multiple files. If "directory", user uploads all files in selected directory. Return type will be list for each file in case of "multiple" or "directory".
type (str): Type of value to be returned by component. "file" returns a temporary file object whose path can be retrieved by file_obj.name, "binary" returns an bytes object.
keepfilename (bool): whether to keep the original filename in the f.name field upon upload. If true, will place 'originalfilename' + a '_' before the unique temporary safe filename string and extension
label (str): component name in interface.
'''
self.file_count = file_count
self.type = type
self.test_input = None
self.keepfilename = keepfilename
super().__init__(label)
def get_template_context(self):
@ -921,7 +924,9 @@ class File(InputComponent):
if is_local_example:
return open(name)
else:
return processing_utils.decode_base64_to_file(data)
if self.keepfilename: filenameprefix=Path(name).stem+'_'
else: filenameprefix=""
return processing_utils.decode_base64_to_file(data, filenameprefix=filenameprefix)
elif self.type == "bytes":
if is_local_example:
with open(name, "rb") as file_data:
@ -930,7 +935,7 @@ class File(InputComponent):
else:
raise ValueError("Unknown type: " + str(self.type) + ". Please choose from: 'file', 'bytes'.")
if self.file_count == "single":
return process_single_file(x[0])
return process_single_file(x)
else:
return [process_single_file(f) for f in x]

View File

@ -80,14 +80,15 @@ def decode_base64_to_binary(encoding):
data = encoding
return base64.b64decode(data), extension
def decode_base64_to_file(encoding, encryption_key=None):
def decode_base64_to_file(encoding, encryption_key=None, filenameprefix=""):
data, extension = decode_base64_to_binary(encoding)
if extension is None:
file_obj = tempfile.NamedTemporaryFile(delete=False)
file_obj = tempfile.NamedTemporaryFile(delete=False, prefix=filenameprefix)
else:
file_obj = tempfile.NamedTemporaryFile(delete=False, suffix="."+extension)
file_obj = tempfile.NamedTemporaryFile(delete=False, prefix=filenameprefix, suffix="."+extension)
if encryption_key is not None:
data = encryptor.encrypt(encryption_key, data)
#print("saving to ", file_obj.name)
file_obj.write(data)
file_obj.flush()
return file_obj