Integrate ask_copy_part into exporters

This introduces some duplication but that'll be refactored later.
This commit is contained in:
Shadab Zafar 2016-07-16 17:48:02 +05:30
parent 8f4d49e22a
commit e3670f9d9b
3 changed files with 23 additions and 27 deletions

View File

@ -194,7 +194,7 @@ def copy_to_clipboard_or_prompt(data):
) )
def flow_format_data(key, scope, flow): def format_flow_data(key, scope, flow):
if key == "u": if key == "u":
data = flow.request.url data = flow.request.url
else: else:
@ -227,12 +227,12 @@ def flow_format_data(key, scope, flow):
return data, False return data, False
def copy_flow(key, scope, flow, writer): def write_flow_data(key, scope, flow, writer):
""" """
key: _c_ontent, _h_eaders+content, _u_rl key: _c_ontent, _h_eaders+content, _u_rl
scope: _a_ll, re_q_uest, re_s_ponse scope: _a_ll, re_q_uest, re_s_ponse
""" """
data, err = flow_format_data(key, scope, flow) data, err = format_flow_data(key, scope, flow)
if err: if err:
signals.status_message.send(message=err) signals.status_message.send(message=err)
@ -250,22 +250,6 @@ def copy_flow(key, scope, flow, writer):
writer(data) writer(data)
def ask_copy_part(scope, flow):
choices = [
("content", "c"),
("headers+content", "h")
]
if scope != "s":
choices.append(("url", "u"))
signals.status_prompt_onekey.send(
prompt = "Copy",
keys = choices,
callback = copy_flow,
args = (scope, flow)
)
def ask_save_body(scope, flow): def ask_save_body(scope, flow):
""" """
Save either the request or the response body to disk. Save either the request or the response body to disk.
@ -315,9 +299,18 @@ def export_to_clip_or_file(key, scope, flow, writer):
'writer' is a function that handles the data 'writer' is a function that handles the data
can be: copy_to_clipboard_or_prompt or ask_save_path can be: copy_to_clipboard_or_prompt or ask_save_path
""" """
for exporter in export.EXPORTERS: for exporter in export.EXPORTERS:
if key == exporter[1]: if key in ["c", "h", "u"]:
writer(exporter[2](flow)) write_flow_data(key, scope, flow, writer)
if key == "r":
writer(export.curl_command(flow))
if key == "p":
writer(export.python_code(flow))
if key == "l":
writer(export.locust_code(flow))
if key == "t":
writer(export.locust_task(flow))
flowcache = utils.LRUCache(800) flowcache = utils.LRUCache(800)

View File

@ -268,7 +268,7 @@ class ConnectionItem(urwid.WidgetWrap):
signals.status_prompt_onekey.send( signals.status_prompt_onekey.send(
self, self,
prompt = "Export to file", prompt = "Export to file",
keys = [(e[0], e[1]) for e in export.EXPORTERS], keys = export.EXPORTERS,
callback = common.export_to_clip_or_file, callback = common.export_to_clip_or_file,
args = ("a", self.flow, common.ask_save_path) args = ("a", self.flow, common.ask_save_path)
) )
@ -276,7 +276,7 @@ class ConnectionItem(urwid.WidgetWrap):
signals.status_prompt_onekey.send( signals.status_prompt_onekey.send(
self, self,
prompt = "Export to clipboard", prompt = "Export to clipboard",
keys = [(e[0], e[1]) for e in export.EXPORTERS], keys = export.EXPORTERS,
callback = common.export_to_clip_or_file, callback = common.export_to_clip_or_file,
args = ("a", self.flow, common.copy_to_clipboard_or_prompt) args = ("a", self.flow, common.copy_to_clipboard_or_prompt)
) )

View File

@ -195,8 +195,11 @@ def locust_task(flow):
EXPORTERS = [ EXPORTERS = [
("as curl command", "c", curl_command), ("content", "c"),
("as python code", "p", python_code), ("headers+content", "h"),
("as locust code", "l", locust_code), ("url", "u"),
("as locust task", "t", locust_task), ("as curl command", "r"),
("as python code", "p"),
("as locust code", "l"),
("as locust task", "t"),
] ]