forked from extern/httpie-cli
289 lines
9.0 KiB
Python
289 lines
9.0 KiB
Python
import argparse
|
|
import os
|
|
import platform
|
|
import sys
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
import requests
|
|
from pygments import __version__ as pygments_version
|
|
from requests import __version__ as requests_version
|
|
|
|
from httpie import __version__ as httpie_version
|
|
from httpie.cli.constants import (
|
|
OUT_REQ_BODY, OUT_REQ_HEAD, OUT_RESP_BODY,
|
|
OUT_RESP_HEAD,
|
|
)
|
|
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.registry import plugin_manager
|
|
from httpie.status import ExitStatus, http_status_to_exit_status
|
|
|
|
|
|
# noinspection PyDefaultArgument
|
|
def main(
|
|
args: List[Union[str, bytes]] = sys.argv,
|
|
env=Environment(),
|
|
) -> ExitStatus:
|
|
"""
|
|
The main function.
|
|
|
|
Pre-process args, handle some special types of invocations,
|
|
and run the main program with error handling.
|
|
|
|
Return exit status code.
|
|
|
|
"""
|
|
program_name, *args = args
|
|
env.program_name = os.path.basename(program_name)
|
|
args = decode_raw_args(args, env.stdin_encoding)
|
|
plugin_manager.load_installed_plugins()
|
|
|
|
from httpie.cli.definition import parser
|
|
|
|
if env.config.default_options:
|
|
args = env.config.default_options + args
|
|
|
|
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,
|
|
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:
|
|
env.stderr.write('\n')
|
|
if include_traceback:
|
|
raise
|
|
exit_status = ExitStatus.ERROR
|
|
else:
|
|
try:
|
|
exit_status = program(
|
|
args=parsed_args,
|
|
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:
|
|
env.stderr.write('\n')
|
|
if include_traceback:
|
|
raise
|
|
exit_status = ExitStatus.ERROR
|
|
except requests.Timeout:
|
|
exit_status = ExitStatus.ERROR_TIMEOUT
|
|
env.log_error(f'Request timed out ({parsed_args.timeout}s).')
|
|
except requests.TooManyRedirects:
|
|
exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS
|
|
env.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}'
|
|
)
|
|
env.log_error(f'{type(e).__name__}: {msg}')
|
|
if include_traceback:
|
|
raise
|
|
exit_status = ExitStatus.ERROR
|
|
|
|
return exit_status
|
|
|
|
|
|
def get_output_options(
|
|
args: argparse.Namespace,
|
|
message: Union[requests.PreparedRequest, requests.Response]
|
|
) -> Tuple[bool, bool]:
|
|
return {
|
|
requests.PreparedRequest: (
|
|
OUT_REQ_HEAD in args.output_options,
|
|
OUT_REQ_BODY in args.output_options,
|
|
),
|
|
requests.Response: (
|
|
OUT_RESP_HEAD in args.output_options,
|
|
OUT_RESP_BODY in args.output_options,
|
|
),
|
|
}[type(message)]
|
|
|
|
|
|
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 print_debug_info(env: Environment):
|
|
env.stderr.writelines([
|
|
f'HTTPie {httpie_version}\n',
|
|
f'Requests {requests_version}\n',
|
|
f'Pygments {pygments_version}\n',
|
|
f'Python {sys.version}\n{sys.executable}\n',
|
|
f'{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
|
|
]
|