codekingpro/portable-devtools

1import argparse2import os3import platform4import sys5import socket6from typing import List, Optional, Union, Callable7 8import requests9from pygments import __version__ as pygments_version10from requests import __version__ as requests_version11 12from . import __version__ as httpie_version13from .cli.constants import OUT_REQ_BODY14from .cli.nested_json import NestedJSONSyntaxError15from .client import collect_messages16from .context import Environment, LogLevel17from .downloads import Downloader18from .models import (19    RequestsMessageKind,20    OutputOptions21)22from .output.models import ProcessingOptions23from .output.writer import write_message, write_stream, write_raw_data, MESSAGE_SEPARATOR_BYTES24from .plugins.registry import plugin_manager25from .status import ExitStatus, http_status_to_exit_status26from .utils import unwrap_context27from .internal.update_warnings import check_updates28from .internal.daemon_runner import is_daemon_mode, run_daemon_task29 30 31# noinspection PyDefaultArgument32def raw_main(33    parser: argparse.ArgumentParser,34    main_program: Callable[[argparse.Namespace, Environment], ExitStatus],35    args: List[Union[str, bytes]] = sys.argv,36    env: Environment = Environment(),37    use_default_options: bool = True,38) -> ExitStatus:39    program_name, *args = args40    env.program_name = os.path.basename(program_name)41    args = decode_raw_args(args, env.stdin_encoding)42 43    if is_daemon_mode(args):44        return run_daemon_task(env, args)45 46    plugin_manager.load_installed_plugins(env.config.plugins_dir)47 48    if use_default_options and env.config.default_options:49        args = env.config.default_options + args50 51    include_debug_info = '--debug' in args52    include_traceback = include_debug_info or '--traceback' in args53 54    def handle_generic_error(e, annotation=None):55        msg = str(e)56        if hasattr(e, 'request'):57            request = e.request58            if hasattr(request, 'url'):59                msg = (60                    f'{msg} while doing a {request.method}'61                    f' request to URL: {request.url}'62                )63        if annotation:64            msg += annotation65        env.log_error(f'{type(e).__name__}: {msg}')66        if include_traceback:67            raise68 69    if include_debug_info:70        print_debug_info(env)71        if args == ['--debug']:72            return ExitStatus.SUCCESS73 74    exit_status = ExitStatus.SUCCESS75 76    try:77        parsed_args = parser.parse_args(78            args=args,79            env=env,80        )81    except NestedJSONSyntaxError as exc:82        env.stderr.write(str(exc) + "\n")83        if include_traceback:84            raise85        exit_status = ExitStatus.ERROR86    except KeyboardInterrupt:87        env.stderr.write('\n')88        if include_traceback:89            raise90        exit_status = ExitStatus.ERROR_CTRL_C91    except SystemExit as e:92        if e.code != ExitStatus.SUCCESS:93            env.stderr.write('\n')94            if include_traceback:95                raise96            exit_status = ExitStatus.ERROR97    else:98        check_updates(env)99        try:100            exit_status = main_program(101                args=parsed_args,102                env=env,103            )104        except KeyboardInterrupt:105            env.stderr.write('\n')106            if include_traceback:107                raise108            exit_status = ExitStatus.ERROR_CTRL_C109        except SystemExit as e:110            if e.code != ExitStatus.SUCCESS:111                env.stderr.write('\n')112                if include_traceback:113                    raise114                exit_status = ExitStatus.ERROR115        except requests.Timeout:116            exit_status = ExitStatus.ERROR_TIMEOUT117            env.log_error(f'Request timed out ({parsed_args.timeout}s).')118        except requests.TooManyRedirects:119            exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS120            env.log_error(121                f'Too many redirects'122                f' (--max-redirects={parsed_args.max_redirects}).'123            )124        except requests.exceptions.ConnectionError as exc:125            annotation = None126            original_exc = unwrap_context(exc)127            if isinstance(original_exc, socket.gaierror):128                if original_exc.errno == socket.EAI_AGAIN:129                    annotation = '\nCouldn’t connect to a DNS server. Please check your connection and try again.'130                elif original_exc.errno == socket.EAI_NONAME:131                    annotation = '\nCouldn’t resolve the given hostname. Please check the URL and try again.'132                propagated_exc = original_exc133            else:134                propagated_exc = exc135 136            handle_generic_error(propagated_exc, annotation=annotation)137            exit_status = ExitStatus.ERROR138        except Exception as e:139            # TODO: Further distinction between expected and unexpected errors.140            handle_generic_error(e)141            exit_status = ExitStatus.ERROR142 143    return exit_status144 145 146def main(147    args: List[Union[str, bytes]] = sys.argv,148    env: Environment = Environment()149) -> ExitStatus:150    """151    The main function.152 153    Pre-process args, handle some special types of invocations,154    and run the main program with error handling.155 156    Return exit status code.157 158    """159 160    from .cli.definition import parser161 162    return raw_main(163        parser=parser,164        main_program=program,165        args=args,166        env=env167    )168 169 170def program(args: argparse.Namespace, env: Environment) -> ExitStatus:171    """172    The main program without error handling.173 174    """175    # TODO: Refactor and drastically simplify, especially so that the separator logic is elsewhere.176    exit_status = ExitStatus.SUCCESS177    downloader = None178    initial_request: Optional[requests.PreparedRequest] = None179    final_response: Optional[requests.Response] = None180    processing_options = ProcessingOptions.from_raw_args(args)181 182    def separate():183        getattr(env.stdout, 'buffer', env.stdout).write(MESSAGE_SEPARATOR_BYTES)184 185    def request_body_read_callback(chunk: bytes):186        should_pipe_to_stdout = bool(187            # Request body output desired188            OUT_REQ_BODY in args.output_options189            # & not `.read()` already pre-request (e.g., for  compression)190            and initial_request191            # & non-EOF chunk192            and chunk193        )194        if should_pipe_to_stdout:195            return write_raw_data(196                env,197                chunk,198                processing_options=processing_options,199                headers=initial_request.headers200            )201 202    try:203        if args.download:204            args.follow = True  # --download implies --follow.205            downloader = Downloader(env, output_file=args.output_file, resume=args.download_resume)206            downloader.pre_request(args.headers)207        messages = collect_messages(env, args=args,208                                    request_body_read_callback=request_body_read_callback)209        force_separator = False210        prev_with_body = False211 212        # Process messages as they’re generated213        for message in messages:214            output_options = OutputOptions.from_message(message, args.output_options)215 216            do_write_body = output_options.body217            if prev_with_body and output_options.any() and (force_separator or not env.stdout_isatty):218                # Separate after a previous message with body, if needed. See test_tokens.py.219                separate()220            force_separator = False221            if output_options.kind is RequestsMessageKind.REQUEST:222                if not initial_request:223                    initial_request = message224                if output_options.body:225                    is_streamed_upload = not isinstance(message.body, (str, bytes))226                    do_write_body = not is_streamed_upload227                    force_separator = is_streamed_upload and env.stdout_isatty228            else:229                final_response = message230                if args.check_status or downloader:231                    exit_status = http_status_to_exit_status(http_status=message.status_code, follow=args.follow)232                    if exit_status != ExitStatus.SUCCESS and (not env.stdout_isatty or args.quiet == 1):233                        env.log_error(f'HTTP {message.raw.status} {message.raw.reason}', level=LogLevel.WARNING)234            write_message(235                requests_message=message,236                env=env,237                output_options=output_options._replace(238                    body=do_write_body239                ),240                processing_options=processing_options241            )242            prev_with_body = output_options.body243 244        # Cleanup245        if force_separator:246            separate()247        if downloader and exit_status == ExitStatus.SUCCESS:248            # Last response body download.249            download_stream, download_to = downloader.start(250                initial_url=initial_request.url,251                final_response=final_response,252            )253            write_stream(stream=download_stream, outfile=download_to, flush=False)254            downloader.finish()255            if downloader.interrupted:256                exit_status = ExitStatus.ERROR257                env.log_error(258                    f'Incomplete download: size={downloader.status.total_size};'259                    f' downloaded={downloader.status.downloaded}'260                )261        return exit_status262 263    finally:264        if downloader and not downloader.finished:265            downloader.failed()266        if args.output_file and args.output_file_specified:267            args.output_file.close()268 269 270def print_debug_info(env: Environment):271    env.stderr.writelines([272        f'HTTPie {httpie_version}\n',273        f'Requests {requests_version}\n',274        f'Pygments {pygments_version}\n',275        f'Python {sys.version}\n{sys.executable}\n',276        f'{platform.system()} {platform.release()}',277    ])278    env.stderr.write('\n\n')279    env.stderr.write(repr(env))280    env.stderr.write('\n\n')281    env.stderr.write(repr(plugin_manager))282    env.stderr.write('\n')283 284 285def decode_raw_args(286    args: List[Union[str, bytes]],287    stdin_encoding: str288) -> List[str]:289    """290    Convert all bytes args to str291    by decoding them using stdin encoding.292 293    """294    return [295        arg.decode(stdin_encoding)296        if type(arg) is bytes else arg297        for arg in args298    ]299