Пример #1
0
    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_no_Content_Length(self, httpbin_both):
     devnull = open(os.devnull, 'w')
     downloader = Downloader(output_file=devnull, progress_file=devnull)
     downloader.start(Response(url=httpbin_both.url + '/'))
     time.sleep(1.1)
     downloader.chunk_downloaded(b'12345')
     downloader.finish()
     assert not downloader.interrupted
Пример #3
0
 def test_download_no_Content_Length(self, httpbin_both):
     devnull = open(os.devnull, 'w')
     downloader = Downloader(output_file=devnull, progress_file=devnull)
     downloader.start(Response(url=httpbin_both.url + '/'))
     time.sleep(1.1)
     downloader.chunk_downloaded(b'12345')
     downloader.finish()
     assert not downloader.interrupted
Пример #4
0
 def test_download_interrupted(self, httpbin_both):
     devnull = open(os.devnull, 'w')
     downloader = Downloader(output_file=devnull, progress_file=devnull)
     downloader.start(
         Response(url=httpbin_both.url + '/', headers={'Content-Length':
                                                       5}))
     downloader.chunk_downloaded(b'1234')
     downloader.finish()
     assert downloader.interrupted
 def test_download_interrupted(self, httpbin_both):
     devnull = open(os.devnull, 'w')
     downloader = Downloader(output_file=devnull, progress_file=devnull)
     downloader.start(Response(
         url=httpbin_both.url + '/',
         headers={'Content-Length': 5}
     ))
     downloader.chunk_downloaded(b'1234')
     downloader.finish()
     assert downloader.interrupted
Пример #6
0
 def test_download_interrupted(self, httpbin_both):
     with open(os.devnull, 'w') as devnull:
         downloader = Downloader(output_file=devnull, progress_file=devnull)
         downloader.start(final_response=Response(
             url=httpbin_both.url + '/', headers={'Content-Length': 5}),
                          initial_url='/')
         downloader.chunk_downloaded(b'1234')
         downloader.finish()
         assert downloader.interrupted
         downloader._progress_reporter.join()
Пример #7
0
 def test_download_no_Content_Length(self, httpbin_both):
     with open(os.devnull, 'w') as devnull:
         downloader = Downloader(output_file=devnull, progress_file=devnull)
         downloader.start(final_response=Response(url=httpbin_both.url +
                                                  '/'),
                          initial_url='/')
         time.sleep(1.1)
         downloader.chunk_downloaded(b'12345')
         downloader.finish()
         assert not downloader.interrupted
         downloader._progress_reporter.join()
Пример #8
0
 def test_download_with_Content_Length(self, httpbin_both):
     with open(os.devnull, 'w') as devnull:
         downloader = Downloader(output_file=devnull, progress_file=devnull)
         downloader.start(Response(
             url=httpbin_both.url + '/',
             headers={'Content-Length': 10}
         ))
         time.sleep(1.1)
         downloader.chunk_downloaded(b'12345')
         time.sleep(1.1)
         downloader.chunk_downloaded(b'12345')
         downloader.finish()
         assert not downloader.interrupted
         downloader._progress_reporter.join()
Пример #9
0
    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()
Пример #10
0
 def test_download_with_Content_Length(self, mock_env, httpbin_both):
     with open(os.devnull, 'w') as devnull:
         downloader = Downloader(mock_env, output_file=devnull)
         downloader.start(
             initial_url='/',
             final_response=Response(
                 url=httpbin_both.url + '/',
                 headers={'Content-Length': 10}
             )
         )
         time.sleep(1.1)
         downloader.chunk_downloaded(b'12345')
         time.sleep(1.1)
         downloader.chunk_downloaded(b'12345')
         downloader.finish()
         assert not downloader.interrupted
Пример #11
0
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()
Пример #12
0
Файл: core.py Проект: co360/TIL
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()
Пример #13
0
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
Пример #14
0
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()
Пример #15
0
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()
Пример #16
0
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()