1import argparse 2import os 3import platform 4import sys 5from typing import List, Optional, Tuple, Union 6 7import requests 8from pygments import __version__ as pygments_version 9from requests import __version__ as requests_version 10 11from . import __version__ as httpie_version 12from .cli.constants import OUT_REQ_BODY, OUT_REQ_HEAD, OUT_RESP_BODY, OUT_RESP_HEAD 13from .client import collect_messages 14from .context import Environment 15from .downloads import Downloader 16from .output.writer import write_message, write_stream, MESSAGE_SEPARATOR_BYTES 17from .plugins.registry import plugin_manager 18from .status import ExitStatus, http_status_to_exit_status 19 20 21# noinspection PyDefaultArgument 22def main(args: List[Union[str, bytes]] = sys.argv, env=Environment()) -> ExitStatus: 23 """ 24 The main function. 25 26 Pre-process args, handle some special types of invocations, 27 and run the main program with error handling. 28 29 Return exit status code. 30 31 """ 32 program_name, *args = args 33 env.program_name = os.path.basename(program_name) 34 args = decode_raw_args(args, env.stdin_encoding) 35 plugin_manager.load_installed_plugins() 36 37 from .cli.definition import parser 38 39 if env.config.default_options: 40 args = env.config.default_options + args 41 42 include_debug_info = '--debug' in args 43 include_traceback = include_debug_info or '--traceback' in args 44 45 if include_debug_info: 46 print_debug_info(env) 47 if args == ['--debug']: 48 return ExitStatus.SUCCESS 49 50 exit_status = ExitStatus.SUCCESS 51 52 try: 53 parsed_args = parser.parse_args( 54 args=args, 55 env=env, 56 ) 57 except KeyboardInterrupt: 58 env.stderr.write('\n') 59 if include_traceback: 60 raise 61 exit_status = ExitStatus.ERROR_CTRL_C 62 except SystemExit as e: 63 if e.code != ExitStatus.SUCCESS: 64 env.stderr.write('\n') 65 if include_traceback: 66 raise 67 exit_status = ExitStatus.ERROR 68 else: 69 try: 70 exit_status = program( 71 args=parsed_args, 72 env=env, 73 ) 74 except KeyboardInterrupt: 75 env.stderr.write('\n') 76 if include_traceback: 77 raise 78 exit_status = ExitStatus.ERROR_CTRL_C 79 except SystemExit as e: 80 if e.code != ExitStatus.SUCCESS: 81 env.stderr.write('\n') 82 if include_traceback: 83 raise 84 exit_status = ExitStatus.ERROR 85 except requests.Timeout: 86 exit_status = ExitStatus.ERROR_TIMEOUT 87 env.log_error(f'Request timed out ({parsed_args.timeout}s).') 88 except requests.TooManyRedirects: 89 exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS 90 env.log_error( 91 f'Too many redirects' 92 f' (--max-redirects={parsed_args.max_redirects}).' 93 ) 94 except Exception as e: 95 # TODO: Further distinction between expected and unexpected errors. 96 msg = str(e) 97 if hasattr(e, 'request'): 98 request = e.request 99 if hasattr(request, 'url'): 100 msg = ( 101 f'{msg} while doing a {request.method}' 102 f' request to URL: {request.url}' 103 ) 104 env.log_error(f'{type(e).__name__}: {msg}') 105 if include_traceback: 106 raise 107 exit_status = ExitStatus.ERROR 108 109 return exit_status 110 111 112def get_output_options( 113 args: argparse.Namespace, 114 message: Union[requests.PreparedRequest, requests.Response] 115) -> Tuple[bool, bool]: 116 return { 117 requests.PreparedRequest: ( 118 OUT_REQ_HEAD in args.output_options, 119 OUT_REQ_BODY in args.output_options, 120 ), 121 requests.Response: ( 122 OUT_RESP_HEAD in args.output_options, 123 OUT_RESP_BODY in args.output_options, 124 ), 125 }[type(message)] 126 127 128def program(args: argparse.Namespace, env: Environment) -> ExitStatus: 129 """ 130 The main program without error handling. 131 132 """ 133 # TODO: Refactor and drastically simplify, especially so that the separator logic is elsewhere. 134 exit_status = ExitStatus.SUCCESS 135 downloader = None 136 initial_request: Optional[requests.PreparedRequest] = None 137 final_response: Optional[requests.Response] = None 138 139 def separate(): 140 getattr(env.stdout, 'buffer', env.stdout).write(MESSAGE_SEPARATOR_BYTES) 141 142 def request_body_read_callback(chunk: bytes): 143 should_pipe_to_stdout = bool( 144 # Request body output desired 145 OUT_REQ_BODY in args.output_options 146 # & not `.read()` already pre-request (e.g., for compression) 147 and initial_request 148 # & non-EOF chunk 149 and chunk 150 ) 151 if should_pipe_to_stdout: 152 msg = requests.PreparedRequest() 153 msg.is_body_upload_chunk = True 154 msg.body = chunk 155 msg.headers = initial_request.headers 156 write_message(requests_message=msg, env=env, args=args, with_body=True, with_headers=False) 157 158 try: 159 if args.download: 160 args.follow = True # --download implies --follow. 161 downloader = Downloader(output_file=args.output_file, progress_file=env.stderr, resume=args.download_resume) 162 downloader.pre_request(args.headers) 163 messages = collect_messages(args=args, config_dir=env.config.directory, 164 request_body_read_callback=request_body_read_callback) 165 force_separator = False 166 prev_with_body = False 167 168 # Process messages as they’re generated 169 for message in messages: 170 is_request = isinstance(message, requests.PreparedRequest) 171 with_headers, with_body = get_output_options(args=args, message=message) 172 do_write_body = with_body 173 if prev_with_body and (with_headers or with_body) and (force_separator or not env.stdout_isatty): 174 # Separate after a previous message with body, if needed. See test_tokens.py. 175 separate() 176 force_separator = False 177 if is_request: 178 if not initial_request: 179 initial_request = message 180 if with_body: 181 is_streamed_upload = not isinstance(message.body, (str, bytes)) 182 do_write_body = not is_streamed_upload 183 force_separator = is_streamed_upload and env.stdout_isatty 184 else: 185 final_response = message 186 if args.check_status or downloader: 187 exit_status = http_status_to_exit_status(http_status=message.status_code, follow=args.follow) 188 if exit_status != ExitStatus.SUCCESS and (not env.stdout_isatty or args.quiet == 1): 189 env.log_error(f'HTTP {message.raw.status} {message.raw.reason}', level='warning') 190 write_message(requests_message=message, env=env, args=args, with_headers=with_headers, 191 with_body=do_write_body) 192 prev_with_body = with_body 193 194 # Cleanup 195 if force_separator: 196 separate() 197 if downloader and exit_status == ExitStatus.SUCCESS: 198 # Last response body download. 199 download_stream, download_to = downloader.start( 200 initial_url=initial_request.url, 201 final_response=final_response, 202 ) 203 write_stream(stream=download_stream, outfile=download_to, flush=False) 204 downloader.finish() 205 if downloader.interrupted: 206 exit_status = ExitStatus.ERROR 207 env.log_error( 208 f'Incomplete download: size={downloader.status.total_size};' 209 f' downloaded={downloader.status.downloaded}' 210 ) 211 return exit_status 212 213 finally: 214 if downloader and not downloader.finished: 215 downloader.failed() 216 if args.output_file and args.output_file_specified: 217 args.output_file.close() 218 219 220def print_debug_info(env: Environment): 221 env.stderr.writelines([ 222 f'HTTPie {httpie_version}\n', 223 f'Requests {requests_version}\n', 224 f'Pygments {pygments_version}\n', 225 f'Python {sys.version}\n{sys.executable}\n', 226 f'{platform.system()} {platform.release()}', 227 ]) 228 env.stderr.write('\n\n') 229 env.stderr.write(repr(env)) 230 env.stderr.write('\n\n') 231 env.stderr.write(repr(plugin_manager)) 232 env.stderr.write('\n') 233 234 235def decode_raw_args( 236 args: List[Union[str, bytes]], 237 stdin_encoding: str 238) -> List[str]: 239 """ 240 Convert all bytes args to str 241 by decoding them using stdin encoding. 242 243 """ 244 return [ 245 arg.decode(stdin_encoding) 246 if type(arg) is bytes else arg 247 for arg in args 248 ] 249