def test_download_output_from_content_disposition(self, httpbin_both):
with tempfile.TemporaryDirectory() as tmp_dirname, open(
os.devnull, 'w') as devnull:
orig_cwd = os.getcwd()
os.chdir(tmp_dirname)
try:
assert not os.path.isfile('filename.bin')
downloader = Downloader(progress_file=devnull)
downloader.start(final_response=Response(
url=httpbin_both.url + '/',
headers={
'Content-Length':
5,
'Content-Disposition':
'attachment; filename="filename.bin"',
}),
initial_url='/')
downloader.chunk_downloaded(b'12345')
downloader.finish()
downloader.failed() # Stop the reporter
assert not downloader.interrupted
downloader._progress_reporter.join()
# TODO: Auto-close the file in that case?
downloader._output_file.close()
assert os.path.isfile('filename.bin')
finally:
os.chdir(orig_cwd)
def test_download_resumed(self, httpbin_both):
with tempfile.TemporaryDirectory() as tmp_dirname:
file = os.path.join(tmp_dirname, 'file.bin')
with open(file, 'a'):
pass
with open(os.devnull, 'w') as devnull, open(file,
'a+b') as output_file:
# Start and interrupt the transfer after 3 bytes written
downloader = Downloader(output_file=output_file,
progress_file=devnull)
downloader.start(final_response=Response(
url=httpbin_both.url + '/', headers={'Content-Length': 5}),
initial_url='/')
downloader.chunk_downloaded(b'123')
downloader.finish()
downloader.failed()
assert downloader.interrupted
downloader._progress_reporter.join()
# Write bytes
with open(file, 'wb') as fh:
fh.write(b'123')
with open(os.devnull, 'w') as devnull, open(file,
'a+b') as output_file:
# Resume the transfer
downloader = Downloader(output_file=output_file,
progress_file=devnull,
resume=True)
# Ensure `pre_request()` is working as expected too
headers = {}
downloader.pre_request(headers)
assert headers['Accept-Encoding'] == 'identity'
assert headers['Range'] == 'bytes=3-'
downloader.start(final_response=Response(
url=httpbin_both.url + '/',
headers={
'Content-Length': 5,
'Content-Range': 'bytes 3-4/5'
},
status_code=PARTIAL_CONTENT),
initial_url='/')
downloader.chunk_downloaded(b'45')
downloader.finish()
downloader._progress_reporter.join()
def program(
args: argparse.Namespace,
env: Environment,
) -> ExitStatus:
"""
The main program without error handling.
"""
exit_status = ExitStatus.SUCCESS
downloader = None
try:
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume)
downloader.pre_request(args.headers)
needs_separator = False
def maybe_separate():
nonlocal needs_separator
if env.stdout_isatty and needs_separator:
needs_separator = False
getattr(env.stdout, 'buffer', env.stdout).write(b'\n\n')
initial_request: Optional[requests.PreparedRequest] = None
final_response: Optional[requests.Response] = None
def request_body_read_callback(chunk: bytes):
should_pipe_to_stdout = (
# Request body output desired
OUT_REQ_BODY in args.output_options
# & not `.read()` already pre-request (e.g., for compression)
and initial_request
# & non-EOF chunk
and chunk)
if should_pipe_to_stdout:
msg = requests.PreparedRequest()
msg.is_body_upload_chunk = True
msg.body = chunk
msg.headers = initial_request.headers
write_message(requests_message=msg,
env=env,
args=args,
with_body=True,
with_headers=False)
messages = collect_messages(
args=args,
config_dir=env.config.directory,
request_body_read_callback=request_body_read_callback)
for message in messages:
maybe_separate()
is_request = isinstance(message, requests.PreparedRequest)
with_headers, with_body = get_output_options(args=args,
message=message)
if is_request:
if not initial_request:
initial_request = message
is_streamed_upload = not isinstance(
message.body, (str, bytes))
if with_body:
with_body = not is_streamed_upload
needs_separator = is_streamed_upload
else:
final_response = message
if args.check_status or downloader:
exit_status = http_status_to_exit_status(
http_status=message.status_code, follow=args.follow)
if (not env.stdout_isatty
and exit_status != ExitStatus.SUCCESS):
env.log_error(
f'HTTP {message.raw.status} {message.raw.reason}',
level='warning')
write_message(
requests_message=message,
env=env,
args=args,
with_headers=with_headers,
with_body=with_body,
)
maybe_separate()
if downloader and exit_status == ExitStatus.SUCCESS:
# Last response body download.
download_stream, download_to = downloader.start(
initial_url=initial_request.url,
final_response=final_response,
)
write_stream(
stream=download_stream,
outfile=download_to,
flush=False,
)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
env.log_error('Incomplete download: size=%d; downloaded=%d' %
(downloader.status.total_size,
downloader.status.downloaded))
return exit_status
finally:
if downloader and not downloader.finished:
downloader.failed()
if (not isinstance(args, list) and args.output_file
and args.output_file_specified):
args.output_file.close()
def program(
args: argparse.Namespace,
env: Environment,
) -> ExitStatus:
"""
The main program without error handling.
"""
exit_status = ExitStatus.SUCCESS
downloader = None
try:
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume)
downloader.pre_request(args.headers)
initial_request = None
final_response = None
for message in collect_messages(args, env.config.directory):
write_message(
requests_message=message,
env=env,
args=args,
)
if isinstance(message, requests.PreparedRequest):
if not initial_request:
initial_request = message
else:
final_response = message
if args.check_status or downloader:
exit_status = http_status_to_exit_status(
http_status=message.status_code, follow=args.follow)
if (not env.stdout_isatty
and exit_status != ExitStatus.SUCCESS):
env.log_error(
f'HTTP {message.raw.status} {message.raw.reason}',
level='warning')
if downloader and exit_status == ExitStatus.SUCCESS:
# Last response body download.
download_stream, download_to = downloader.start(
initial_url=initial_request.url,
final_response=final_response,
)
write_stream(
stream=download_stream,
outfile=download_to,
flush=False,
)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
env.log_error('Incomplete download: size=%d; downloaded=%d' %
(downloader.status.total_size,
downloader.status.downloaded))
return exit_status
finally:
if downloader and not downloader.finished:
downloader.failed()
if (not isinstance(args, list) and args.output_file
and args.output_file_specified):
args.output_file.close()
def main(args=sys.argv[1:], env=Environment(), error=None):
"""Run the main program and write the output to ``env.stdout``.
Return exit status code.
"""
args = decode_args(args, env.stdin_encoding)
plugin_manager.load_installed_plugins()
from httpie.cli import parser
if env.config.default_options:
args = env.config.default_options + args
def _error(msg, *args, **kwargs):
msg = msg % args
level = kwargs.get('level', 'error')
env.stderr.write('\nhttp: %s: %s\n' % (level, msg))
if error is None:
error = _error
debug = '--debug' in args
traceback = debug or '--traceback' in args
exit_status = ExitStatus.OK
if debug:
print_debug_info(env)
if args == ['--debug']:
return exit_status
downloader = None
try:
args = parser.parse_args(args=args, env=env)
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(
output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume
)
downloader.pre_request(args.headers)
last_response = get_response(args, config_dir=env.config.directory)
if args.show_redirects:
responses = last_response.history + [last_response]
else:
responses = [last_response]
for response in responses:
if args.check_status or downloader:
exit_status = get_exit_status(
http_status=response.status_code,
follow=args.follow
)
if not env.stdout_isatty and exit_status != ExitStatus.OK:
error('HTTP %s %s',
response.raw.status,
response.raw.reason,
level='warning')
write_stream_kwargs = {
'stream': build_output_stream(
args=args,
env=env,
request=response.request,
response=response,
),
# NOTE: `env.stdout` will in fact be `stderr` with `--download`
'outfile': env.stdout,
'flush': env.stdout_isatty or args.stream
}
try:
if env.is_windows and is_py3 and 'colors' in args.prettify:
write_stream_with_colors_win_py3(**write_stream_kwargs)
else:
write_stream(**write_stream_kwargs)
except IOError as e:
if not traceback and e.errno == errno.EPIPE:
# Ignore broken pipes unless --traceback.
env.stderr.write('\n')
else:
raise
if downloader and exit_status == ExitStatus.OK:
# Last response body download.
download_stream, download_to = downloader.start(last_response)
write_stream(
stream=download_stream,
outfile=download_to,
flush=False,
)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
error('Incomplete download: size=%d; downloaded=%d' % (
downloader.status.total_size,
downloader.status.downloaded
))
except KeyboardInterrupt:
if traceback:
raise
env.stderr.write('\n')
exit_status = ExitStatus.ERROR
except SystemExit as e:
if e.code != ExitStatus.OK:
if traceback:
raise
env.stderr.write('\n')
exit_status = ExitStatus.ERROR
except requests.Timeout:
exit_status = ExitStatus.ERROR_TIMEOUT
error('Request timed out (%ss).', args.timeout)
except requests.TooManyRedirects:
exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS
error('Too many redirects (--max-redirects=%s).', args.max_redirects)
except Exception as e:
# TODO: Better distinction between expected and unexpected errors.
if traceback:
raise
msg = str(e)
if hasattr(e, 'request'):
request = e.request
if hasattr(request, 'url'):
msg += ' while doing %s request to URL: %s' % (
request.method, request.url)
error('%s: %s', type(e).__name__, msg)
exit_status = ExitStatus.ERROR
finally:
if downloader and not downloader.finished:
downloader.failed()
return exit_status
def program(args, env, log_error):
"""
The main program without error handling
:param args: parsed args (argparse.Namespace)
:type env: Environment
:param log_error: error log function
:return: status code
"""
exit_status = ExitStatus.OK
downloader = None
show_traceback = args.debug or args.traceback
try:
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume)
downloader.pre_request(args.headers)
final_response = get_response(args, config_dir=env.config.directory)
if args.all:
responses = final_response.history + [final_response]
else:
responses = [final_response]
for response in responses:
if args.check_status or downloader:
exit_status = get_exit_status(http_status=response.status_code,
follow=args.follow)
if not env.stdout_isatty and exit_status != ExitStatus.OK:
log_error('HTTP %s %s',
response.raw.status,
response.raw.reason,
level='warning')
write_stream_kwargs = {
'stream':
build_output_stream(
args=args,
env=env,
request=response.request,
response=response,
output_options=(args.output_options
if response is final_response else
args.output_options_others)),
# NOTE: `env.stdout` will in fact be `stderr` with `--download`
'outfile':
env.stdout,
'flush':
env.stdout_isatty or args.stream
}
try:
if env.is_windows and is_py3 and 'colors' in args.prettify:
write_stream_with_colors_win_py3(**write_stream_kwargs)
else:
write_stream(**write_stream_kwargs)
except IOError as e:
if not show_traceback and e.errno == errno.EPIPE:
# Ignore broken pipes unless --traceback.
env.stderr.write('\n')
else:
raise
if downloader and exit_status == ExitStatus.OK:
# Last response body download.
download_stream, download_to = downloader.start(final_response)
write_stream(
stream=download_stream,
outfile=download_to,
flush=False,
)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
log_error('Incomplete download: size=%d; downloaded=%d' %
(downloader.status.total_size,
downloader.status.downloaded))
return exit_status
finally:
if downloader and not downloader.finished:
downloader.failed()
if (not isinstance(args, list) and args.output_file
and args.output_file_specified):
args.output_file.close()
def program(args, env, log_error):
"""
The main program without error handling
:param args: parsed args (argparse.Namespace)
:type env: Environment
:param log_error: error log function
:return: status code
"""
exit_status = ExitStatus.OK
downloader = None
show_traceback = args.debug or args.traceback
try:
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(
output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume
)
downloader.pre_request(args.headers)
final_response = get_response(args, config_dir=env.config.directory)
if args.all:
responses = final_response.history + [final_response]
else:
responses = [final_response]
for response in responses:
if args.check_status or downloader:
exit_status = get_exit_status(
http_status=response.status_code,
follow=args.follow
)
if not env.stdout_isatty and exit_status != ExitStatus.OK:
log_error(
'HTTP %s %s', response.raw.status, response.raw.reason,
level='warning'
)
write_stream_kwargs = {
'stream': build_output_stream(
args=args,
env=env,
request=response.request,
response=response,
output_options=(
args.output_options
if response is final_response
else args.output_options_history
)
),
# NOTE: `env.stdout` will in fact be `stderr` with `--download`
'outfile': env.stdout,
'flush': env.stdout_isatty or args.stream
}
try:
if env.is_windows and is_py3 and 'colors' in args.prettify:
write_stream_with_colors_win_py3(**write_stream_kwargs)
else:
write_stream(**write_stream_kwargs)
except IOError as e:
if not show_traceback and e.errno == errno.EPIPE:
# Ignore broken pipes unless --traceback.
env.stderr.write('\n')
else:
raise
if downloader and exit_status == ExitStatus.OK:
# Last response body download.
download_stream, download_to = downloader.start(final_response)
write_stream(
stream=download_stream,
outfile=download_to,
flush=False,
)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
log_error('Incomplete download: size=%d; downloaded=%d' % (
downloader.status.total_size,
downloader.status.downloaded
))
return exit_status
finally:
if downloader and not downloader.finished:
downloader.failed()
if (not isinstance(args, list) and args.output_file and
args.output_file_specified):
args.output_file.close()
def program(args: argparse.Namespace, env: Environment) -> ExitStatus:
"""
The main program without error handling.
"""
# TODO: Refactor and drastically simplify, especially so that the separator logic is elsewhere.
exit_status = ExitStatus.SUCCESS
downloader = None
initial_request: Optional[requests.PreparedRequest] = None
final_response: Optional[requests.Response] = None
def separate():
getattr(env.stdout, 'buffer',
env.stdout).write(MESSAGE_SEPARATOR_BYTES)
def request_body_read_callback(chunk: bytes):
should_pipe_to_stdout = bool(
# Request body output desired
OUT_REQ_BODY in args.output_options
# & not `.read()` already pre-request (e.g., for compression)
and initial_request
# & non-EOF chunk
and chunk)
if should_pipe_to_stdout:
msg = requests.PreparedRequest()
msg.is_body_upload_chunk = True
msg.body = chunk
msg.headers = initial_request.headers
write_message(requests_message=msg,
env=env,
args=args,
with_body=True,
with_headers=False)
try:
if args.download:
args.follow = True # --download implies --follow.
downloader = Downloader(output_file=args.output_file,
progress_file=env.stderr,
resume=args.download_resume)
downloader.pre_request(args.headers)
messages = collect_messages(
args=args,
config_dir=env.config.directory,
request_body_read_callback=request_body_read_callback)
force_separator = False
prev_with_body = False
if args.output_format_form == "RAW":
# Process messages as they’re generated
for message in messages:
is_request = isinstance(message, requests.PreparedRequest)
with_headers, with_body = get_output_options(args=args,
message=message)
do_write_body = with_body
if prev_with_body and (with_headers or with_body) and (
force_separator or not env.stdout_isatty):
# Separate after a previous message with body, if needed. See test_tokens.py.
separate()
force_separator = False
if is_request:
if not initial_request:
initial_request = message
is_streamed_upload = not isinstance(
message.body, (str, bytes))
if with_body:
do_write_body = not is_streamed_upload
force_separator = is_streamed_upload and env.stdout_isatty
else:
final_response = message
if args.check_status or downloader:
exit_status = http_status_to_exit_status(
http_status=message.status_code,
follow=args.follow)
if exit_status != ExitStatus.SUCCESS and (
not env.stdout_isatty or args.quiet):
env.log_error(
f'HTTP {message.raw.status} {message.raw.reason}',
level='warning')
write_message(requests_message=message,
env=env,
args=args,
with_headers=with_headers,
with_body=do_write_body)
prev_with_body = with_body
else:
all_messages_together = []
for message in messages:
is_request = isinstance(message, requests.PreparedRequest)
with_headers, with_body = get_output_options(args=args,
message=message)
#force_separator = False
if is_request:
with_headers_req = with_headers
with_body_req = with_body
if not initial_request:
initial_request = message
is_streamed_upload = not isinstance(
message.body, (str, bytes))
if with_body:
with_body_req = not is_streamed_upload
#force_separator = is_streamed_upload and env.stdout_isatty
else:
with_headers_res = with_headers
with_body_res = with_body
final_response = message
if args.check_status or downloader:
exit_status = http_status_to_exit_status(
http_status=message.status_code,
follow=args.follow)
if exit_status != ExitStatus.SUCCESS and (
not env.stdout_isatty or args.quiet):
env.log_error(
f'HTTP {message.raw.status} {message.raw.reason}',
level='warning')
all_messages_together.append(message)
write_message_json(requests_message=all_messages_together,
env=env,
args=args,
with_headers_req=with_headers_req,
with_body_req=with_body_req,
with_headers_res=with_headers_res,
with_body_res=with_body_res)
prev_with_body = with_body
# Cleanup
if force_separator:
separate()
if downloader and exit_status == ExitStatus.SUCCESS:
# Last response body download.
download_stream, download_to = downloader.start(
initial_url=initial_request.url,
final_response=final_response,
)
write_stream(stream=download_stream,
outfile=download_to,
flush=False)
downloader.finish()
if downloader.interrupted:
exit_status = ExitStatus.ERROR
env.log_error('Incomplete download: size=%d; downloaded=%d' %
(downloader.status.total_size,
downloader.status.downloaded))
return exit_status
finally:
if downloader and not downloader.finished:
downloader.failed()
if not isinstance(
args,
list) and args.output_file and args.output_file_specified:
args.output_file.close()