Skip to content

Commit

Permalink
Refactor FileOperator (#787)
Browse files Browse the repository at this point in the history
  • Loading branch information
jopohl authored Jul 8, 2020
1 parent 22f3cf6 commit 716384d
Show file tree
Hide file tree
Showing 8 changed files with 109 additions and 104 deletions.
2 changes: 1 addition & 1 deletion src/urh/controller/CompareFrameController.py
Original file line number Diff line number Diff line change
Expand Up @@ -857,7 +857,7 @@ def save_protocol(self):
break

text = "protocol"
filename = FileOperator.get_save_file_name("{0}.proto.xml".format(text), caption="Save protocol")
filename = FileOperator.ask_save_file_name("{0}.proto.xml".format(text), caption="Save protocol")

if not filename:
return
Expand Down
4 changes: 2 additions & 2 deletions src/urh/controller/GeneratorTabController.py
Original file line number Diff line number Diff line change
Expand Up @@ -384,7 +384,7 @@ def generate_file(self):
except Exception as e:
logger.exception(e)
sample_rate = 1e6
FileOperator.save_data_dialog("generated", modulated_samples, sample_rate=sample_rate, parent=self)
FileOperator.ask_signal_file_name_and_save("generated", modulated_samples, sample_rate=sample_rate, parent=self)
except Exception as e:
Errors.exception(e)
self.unsetCursor()
Expand Down Expand Up @@ -608,7 +608,7 @@ def on_btn_send_clicked(self):

@pyqtSlot()
def on_btn_save_clicked(self):
filename = FileOperator.get_save_file_name("profile.fuzz.xml", caption="Save fuzz profile")
filename = FileOperator.ask_save_file_name("profile.fuzz.xml", caption="Save fuzzing profile")
if filename:
self.table_model.protocol.to_xml_file(filename,
decoders=self.project_manager.decodings,
Expand Down
2 changes: 1 addition & 1 deletion src/urh/controller/SimulatorTabController.py
Original file line number Diff line number Diff line change
Expand Up @@ -507,7 +507,7 @@ def refresh_tree(self):

@pyqtSlot()
def on_btn_save_clicked(self):
filename = FileOperator.get_save_file_name(initial_name="myprofile.sim.xml", caption="Save simulator profile")
filename = FileOperator.ask_save_file_name(initial_name="myprofile.sim.xml", caption="Save simulator profile")
if filename:
self.save_simulator_file(filename)

Expand Down
4 changes: 2 additions & 2 deletions src/urh/controller/dialogs/ReceiveDialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,8 @@ def on_save_clicked(self):

initial_name = initial_name.replace(Formatter.local_decimal_seperator(), "_").replace("_000", "")

filename = FileOperator.save_data_dialog(initial_name, data,
sample_rate=dev.sample_rate, parent=self)
filename = FileOperator.ask_signal_file_name_and_save(initial_name, data,
sample_rate=dev.sample_rate, parent=self)
self.already_saved = True
if filename is not None and filename not in self.recorded_files:
self.recorded_files.append(filename)
2 changes: 1 addition & 1 deletion src/urh/controller/dialogs/SendDialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def init_device(self):

@pyqtSlot()
def on_graphics_view_save_as_clicked(self):
filename = FileOperator.get_save_file_name("signal.complex")
filename = FileOperator.ask_save_file_name("signal.complex")
if filename:
try:
try:
Expand Down
4 changes: 2 additions & 2 deletions src/urh/controller/dialogs/SimulatorDialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -418,8 +418,8 @@ def on_btn_save_rx_clicked(self):
rx_device = self.simulator.sniffer.rcv_device
if isinstance(rx_device.data, np.ndarray) or isinstance(rx_device.data, IQArray):
data = IQArray(rx_device.data[:rx_device.current_index])
filename = FileOperator.save_data_dialog("simulation_capture", data,
sample_rate=rx_device.sample_rate, parent=self)
filename = FileOperator.ask_signal_file_name_and_save("simulation_capture", data,
sample_rate=rx_device.sample_rate, parent=self)
if filename:
data.tofile(filename)
self.rx_file_saved.emit(filename)
Expand Down
8 changes: 4 additions & 4 deletions src/urh/controller/widgets/SignalFrame.py
Original file line number Diff line number Diff line change
Expand Up @@ -448,8 +448,8 @@ def save_signal(self):

def save_signal_as(self):
try:
FileOperator.save_data_dialog(self.signal.name, self.signal.iq_array, self.signal.sample_rate,
self.signal.wav_mode)
FileOperator.ask_signal_file_name_and_save(self.signal.name, self.signal.iq_array, self.signal.sample_rate,
self.signal.wav_mode)
except Exception as e:
Errors.exception(e)

Expand All @@ -460,7 +460,7 @@ def export_demodulated(self):
logger.exception(e)
initial_name = "demodulated.complex"

filename = FileOperator.get_save_file_name(initial_name)
filename = FileOperator.ask_save_file_name(initial_name)
if filename:
try:
self.setCursor(Qt.WaitCursor)
Expand Down Expand Up @@ -1314,7 +1314,7 @@ def on_export_fta_wanted(self):
logger.exception(e)
initial_name = "spectrogram.ft"

filename = FileOperator.get_save_file_name(initial_name, caption="Export spectrogram")
filename = FileOperator.ask_save_file_name(initial_name, caption="Export spectrogram")
if not filename:
return
QApplication.setOverrideCursor(Qt.WaitCursor)
Expand Down
187 changes: 96 additions & 91 deletions src/urh/util/FileOperator.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,28 +8,53 @@
from PyQt5.QtCore import QDir
from PyQt5.QtWidgets import QFileDialog, QMessageBox

from urh.models.FileIconProvider import FileIconProvider
from urh.signalprocessing.IQArray import IQArray

VIEW_TYPES = ["Bits", "Hex", "ASCII"]

archives = {}
""":type: dict of [str, str]
:param: archives[extracted_filename] = filename"""

RECENT_PATH = QDir.homePath()

EXT = {np.int8: ".complex16s", np.uint8: ".complex16u", np.int16: ".complex32s", np.uint16: ".complex32u",
np.float32: ".complex", np.complex64: ".complex"}
FILTER = {np.int8: "Complex16 signed (*.complex16s *.cs8)", np.uint8: "Complex16 unsigned (*.complex16u *.cu8)",
np.uint16: "Complex32 unsigned (*.complex32u *.cu16)", np.int16: "Complex32 signed (*.complex32s *.cs16)",
np.float32: "Complex (*.complex)", np.complex64: "Complex (*.complex)"}
SIGNAL_FILE_EXTENSIONS_BY_TYPE = {
np.int8: ".complex16s",
np.uint8: ".complex16u",
np.int16: ".complex32s",
np.uint16: ".complex32u",
np.float32: ".complex",
np.complex64: ".complex"
}

SIGNAL_NAME_FILTERS_BY_TYPE = {
np.int8: "Complex16 signed (*.complex16s *.cs8)",
np.uint8: "Complex16 unsigned (*.complex16u *.cu8)",
np.uint16: "Complex32 unsigned (*.complex32u *.cu16)",
np.int16: "Complex32 signed (*.complex32s *.cs16)",
np.float32: "Complex (*.complex)",
np.complex64: "Complex (*.complex)"
}

EVERYTHING_FILE_FILTER = "All Files (*)"

SIGNAL_NAME_FILTERS = list(sorted(set(SIGNAL_NAME_FILTERS_BY_TYPE.values())))

COMPRESSED_COMPLEX_FILE_FILTER = "Compressed Complex File (*.coco)"
WAV_FILE_FILTER = "Waveform Audio File Format (*.wav *.wave)"
PROTOCOL_FILE_FILTER = "Protocol (*.proto.xml *.proto)"
BINARY_PROTOCOL_FILE_FILTER = "Binary Protocol (*.bin)"
PLAIN_BITS_FILE_FILTER = "Plain Bits (*.txt)"
FUZZING_FILE_FILTER = "Fuzzing Profile (*.fuzz.xml *.fuzz)"
SIMULATOR_FILE_FILTER = "Simulator Profile (*.sim.xml *.sim)"
TAR_FILE_FILTER = "Tar Archive (*.tar *.tar.gz *.tar.bz2)"
ZIP_FILE_FILTER = "Zip Archive (*.zip)"


def __get__name_filter_for_signals() -> str:
return ";;".join([EVERYTHING_FILE_FILTER] + SIGNAL_NAME_FILTERS + [COMPRESSED_COMPLEX_FILE_FILTER, WAV_FILE_FILTER])


def get_open_dialog(directory_mode=False, parent=None, name_filter="full") -> QFileDialog:
fip = FileIconProvider()
dialog = QFileDialog(parent=parent, directory=RECENT_PATH)
dialog.setIconProvider(fip)

if directory_mode:
dialog.setFileMode(QFileDialog.Directory)
Expand All @@ -38,98 +63,43 @@ def get_open_dialog(directory_mode=False, parent=None, name_filter="full") -> QF
dialog.setFileMode(QFileDialog.ExistingFiles)
dialog.setWindowTitle("Open Files")
if name_filter == "full":
name_filter = "All Files (*);;" \
"Complex (*.complex);;" \
"Complex16 unsigned (*.complex16u *.cu8);;" \
"Complex16 signed (*.complex16s *.cs8);;" \
"Complex32 unsigned (*.complex32u *.cu16);;" \
"Complex32 signed (*.complex32s *.cs16);;" \
"WAV (*.wav);;" \
"Protocols (*.proto.xml *.proto);;" \
"Binary Protocols (*.bin);;" \
"Fuzzing Profiles (*.fuzz.xml *.fuzz);;" \
"Simulator (*.sim.xml *.sim)" \
"Plain Bits (*.txt);;" \
"Tar Archives (*.tar *.tar.gz *.tar.bz2);;" \
"Zip Archives (*.zip)"
name_filter = __get__name_filter_for_signals() + ";;" \
+ ";;".join([PROTOCOL_FILE_FILTER, BINARY_PROTOCOL_FILE_FILTER, PLAIN_BITS_FILE_FILTER,
FUZZING_FILE_FILTER, SIMULATOR_FILE_FILTER, TAR_FILE_FILTER, ZIP_FILE_FILTER])
elif name_filter == "signals_only":
name_filter = __get__name_filter_for_signals()
elif name_filter == "proto":
name_filter = "Protocols (*.proto.xml *.proto);; Binary Protocols (*.bin)"
name_filter = ";;".join([PROTOCOL_FILE_FILTER, BINARY_PROTOCOL_FILE_FILTER])
elif name_filter == "fuzz":
name_filter = "Fuzzprofiles (*.fuzz.xml *.fuzz)"
name_filter = FUZZING_FILE_FILTER
elif name_filter == "simulator":
name_filter = "Simulator (*.sim.xml *.sim)"
name_filter = SIMULATOR_FILE_FILTER

dialog.setNameFilter(name_filter)

dialog.setOptions(QFileDialog.DontResolveSymlinks)
dialog.setViewMode(QFileDialog.Detail)

return dialog


def uncompress_archives(file_names, temp_dir):
"""
Extract each archive from the list of filenames.
Normal files stay untouched.
Add all files to the Recent Files.
:type file_names: list of str
:type temp_dir: str
:rtype: list of str
"""
result = []
for filename in file_names:
if filename.endswith(".tar") or filename.endswith(".tar.gz") or filename.endswith(".tar.bz2"):
obj = tarfile.open(filename, "r")
extracted_file_names = []
for j, member in enumerate(obj.getmembers()):
obj.extract(member, temp_dir)
extracted_filename = os.path.join(temp_dir, obj.getnames()[j])
extracted_file_names.append(extracted_filename)
archives[extracted_filename] = filename
result.extend(extracted_file_names[:])
elif filename.endswith(".zip"):
obj = zipfile.ZipFile(filename)
extracted_file_names = []
for j, info in enumerate(obj.infolist()):
obj.extract(info, path=temp_dir)
extracted_filename = os.path.join(temp_dir, obj.namelist()[j])
extracted_file_names.append(extracted_filename)
archives[extracted_filename] = filename
result.extend(extracted_file_names[:])
else:
result.append(filename)

return result


def get_save_file_name(initial_name: str, wav_only=False, caption="Save signal", selected_name_filter=None):
def ask_save_file_name(initial_name: str, caption="Save signal", selected_name_filter=None):
global RECENT_PATH
if caption == "Save signal":
name_filter = "Complex (*.complex);;" \
"Complex16 unsigned (*.complex16u *.cu8);;" \
"Complex16 signed (*.complex16s *.cs8);;" \
"Complex32 unsigned (*.complex32u *.cu16);;" \
"Complex32 signed (*.complex32s *.cs16);;" \
"Complex compressed (*.coco);;" \
"WAV (*.wav);;" \
"All Files (*)"
if wav_only:
name_filter = "WAV (*.wav);;All Files (*)"
elif caption == "Save fuzz profile":
name_filter = "Fuzzing Profile (*.fuzz.xml *.fuzz);;All Files (*)"
name_filter = __get__name_filter_for_signals()
elif caption == "Save fuzzing profile":
name_filter = FUZZING_FILE_FILTER
elif caption == "Save encoding":
name_filter = ""
elif caption == "Save simulator profile":
name_filter = "Simulator (*.sim.xml *.sim);;All Files (*)"
name_filter = SIMULATOR_FILE_FILTER
elif caption == "Export spectrogram":
name_filter = "Frequency Time (*.ft);;Frequency Time Amplitude (*.fta)"
elif caption == "Save protocol":
name_filter = ";;".join([PROTOCOL_FILE_FILTER, BINARY_PROTOCOL_FILE_FILTER])
else:
name_filter = "Protocols (*.proto.xml *.proto);;Binary Protocol (*.bin);;All Files (*)"
name_filter = EVERYTHING_FILE_FILTER

filename = None
dialog = QFileDialog(directory=RECENT_PATH, caption=caption, filter=name_filter)
dialog.setFileMode(QFileDialog.AnyFile)
dialog.setViewMode(QFileDialog.Detail)
dialog.setLabelText(QFileDialog.Accept, "Save")
dialog.setAcceptMode(QFileDialog.AcceptSave)

Expand All @@ -147,23 +117,23 @@ def get_save_file_name(initial_name: str, wav_only=False, caption="Save signal",
return filename


def save_data_dialog(signal_name: str, data, sample_rate=1e6, wav_only=False, parent=None) -> str:
def ask_signal_file_name_and_save(signal_name: str, data, sample_rate=1e6, wav_only=False, parent=None) -> str:
if wav_only:
if not signal_name.endswith(".wav"):
if not signal_name.endswith(".wav") and not signal_name.endswith(".wave"):
signal_name += ".wav"
name_filter = "WAV (*.wav)"
selected_name_filter = WAV_FILE_FILTER
else:
if not any(signal_name.endswith(e) for e in FILTER.values()):
if not any(signal_name.endswith(e) for e in SIGNAL_NAME_FILTERS_BY_TYPE.values()):
try:
dtype = next(d for d in EXT.keys() if d == data.dtype)
signal_name += EXT[dtype]
name_filter = FILTER[dtype]
dtype = next(d for d in SIGNAL_FILE_EXTENSIONS_BY_TYPE.keys() if d == data.dtype)
signal_name += SIGNAL_FILE_EXTENSIONS_BY_TYPE[dtype]
selected_name_filter = SIGNAL_NAME_FILTERS_BY_TYPE[dtype]
except StopIteration:
name_filter = None
selected_name_filter = None
else:
name_filter = None
selected_name_filter = None

filename = get_save_file_name(signal_name, wav_only, selected_name_filter=name_filter)
filename = ask_save_file_name(signal_name, selected_name_filter=selected_name_filter)

if filename:
try:
Expand Down Expand Up @@ -232,6 +202,41 @@ def rewrite_tar(tar_name: str):
shutil.rmtree(tempdir)


def uncompress_archives(file_names, temp_dir):
"""
Extract each archive from the list of filenames.
Normal files stay untouched.
Add all files to the Recent Files.
:type file_names: list of str
:type temp_dir: str
:rtype: list of str
"""
result = []
for filename in file_names:
if filename.endswith(".tar") or filename.endswith(".tar.gz") or filename.endswith(".tar.bz2"):
obj = tarfile.open(filename, "r")
extracted_file_names = []
for j, member in enumerate(obj.getmembers()):
obj.extract(member, temp_dir)
extracted_filename = os.path.join(temp_dir, obj.getnames()[j])
extracted_file_names.append(extracted_filename)
archives[extracted_filename] = filename
result.extend(extracted_file_names[:])
elif filename.endswith(".zip"):
obj = zipfile.ZipFile(filename)
extracted_file_names = []
for j, info in enumerate(obj.infolist()):
obj.extract(info, path=temp_dir)
extracted_filename = os.path.join(temp_dir, obj.namelist()[j])
extracted_file_names.append(extracted_filename)
archives[extracted_filename] = filename
result.extend(extracted_file_names[:])
else:
result.append(filename)

return result


def get_directory():
directory = QFileDialog.getExistingDirectory(None, "Choose Directory", QDir.homePath(),
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
Expand Down

0 comments on commit 716384d

Please sign in to comment.