httpie-cli/httpie/core.py

238 lines
7.2 KiB
Python
Raw Normal View History

import argparse
2016-03-04 18:42:13 +01:00
import platform
import sys
from typing import Callable, List, Union
2012-07-26 00:26:23 +02:00
import requests
2012-08-07 14:50:51 +02:00
from pygments import __version__ as pygments_version
from requests import __version__ as requests_version
2012-07-26 00:26:23 +02:00
from httpie import ExitStatus, __version__ as httpie_version
from httpie.client import collect_messages
from httpie.context import Environment
from httpie.downloads import Downloader
from httpie.output.writer import write_message, write_stream
from httpie.plugins import plugin_manager
def main(
args: List[Union[str, bytes]] = sys.argv,
env=Environment(),
custom_log_error: Callable = None
) -> ExitStatus:
"""
The main function.
2016-03-01 14:37:26 +01:00
Pre-process args, handle some special types of invocations,
and run the main program with error handling.
Return exit status code.
"""
args = decode_raw_args(args, env.stdin_encoding)
program_name, *args = args
plugin_manager.load_installed_plugins()
def log_error(msg, level='error'):
2016-03-01 14:27:26 +01:00
assert level in ['error', 'warning']
env.stderr.write(f'\n{program_name}: {level}: {msg}\n')
2019-08-31 15:17:10 +02:00
from httpie.cli.definition import parser
if env.config.default_options:
args = env.config.default_options + args
if custom_log_error:
log_error = custom_log_error
include_debug_info = '--debug' in args
include_traceback = include_debug_info or '--traceback' in args
if include_debug_info:
print_debug_info(env)
if args == ['--debug']:
return ExitStatus.SUCCESS
exit_status = ExitStatus.SUCCESS
try:
parsed_args = parser.parse_args(
args=args,
program_name=program_name,
env=env,
)
except KeyboardInterrupt:
env.stderr.write('\n')
if include_traceback:
raise
exit_status = ExitStatus.ERROR_CTRL_C
except SystemExit as e:
if e.code != ExitStatus.SUCCESS.value:
env.stderr.write('\n')
if include_traceback:
raise
exit_status = ExitStatus.ERROR
else:
try:
exit_status = program(
args=parsed_args,
env=env,
log_error=log_error,
)
except KeyboardInterrupt:
env.stderr.write('\n')
if include_traceback:
raise
exit_status = ExitStatus.ERROR_CTRL_C
except SystemExit as e:
if e.code != ExitStatus.SUCCESS.value:
env.stderr.write('\n')
if include_traceback:
raise
exit_status = ExitStatus.ERROR
except requests.Timeout:
exit_status = ExitStatus.ERROR_TIMEOUT
log_error(f'Request timed out ({parsed_args.timeout}s).')
except requests.TooManyRedirects:
exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS
log_error(
f'Too many redirects'
f' (--max-redirects=parsed_args.max_redirects).'
)
except Exception as e:
# TODO: Further distinction between expected and unexpected errors.
msg = str(e)
if hasattr(e, 'request'):
request = e.request
if hasattr(request, 'url'):
msg = (
f'{msg} while doing a {request.method}'
f' request to URL: {request.url}'
)
log_error(f'{type(e).__name__}: {msg}')
if include_traceback:
raise
exit_status = ExitStatus.ERROR
2012-12-11 12:54:34 +01:00
return exit_status
def program(
args: argparse.Namespace,
env: Environment,
log_error: Callable
) -> 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 = get_exit_status(
http_status=message.status_code,
follow=args.follow
)
if not env.stdout_isatty and exit_status != ExitStatus.SUCCESS:
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
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 get_exit_status(http_status: int, follow=False) -> ExitStatus:
"""Translate HTTP status code to exit status code."""
if 300 <= http_status <= 399 and not follow:
# Redirect
return ExitStatus.ERROR_HTTP_3XX
elif 400 <= http_status <= 499:
# Client Error
return ExitStatus.ERROR_HTTP_4XX
elif 500 <= http_status <= 599:
# Server Error
return ExitStatus.ERROR_HTTP_5XX
else:
return ExitStatus.SUCCESS
def print_debug_info(env: Environment):
env.stderr.writelines([
'HTTPie %s\n' % httpie_version,
'Requests %s\n' % requests_version,
'Pygments %s\n' % pygments_version,
'Python %s\n%s\n' % (sys.version, sys.executable),
'%s %s' % (platform.system(), platform.release()),
])
env.stderr.write('\n\n')
env.stderr.write(repr(env))
env.stderr.write('\n')
def decode_raw_args(
args: List[Union[str, bytes]],
stdin_encoding: str
) -> List[str]:
"""
Convert all bytes args to str
by decoding them using stdin encoding.
"""
return [
arg.decode(stdin_encoding)
if type(arg) == bytes else arg
for arg in args
]