Add one-by-one processing of each HTTP request or response and --offline

httpie/client.py

@@ -5,14 +5,16 @@ import sys
 import zlib
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Iterable, Union
 
 import requests
 from requests.adapters import HTTPAdapter
 
-from httpie import __version__, sessions
+from httpie import __version__
 from httpie.cli.constants import SSL_VERSION_ARG_MAPPING
 from httpie.cli.dicts import RequestHeadersDict
 from httpie.plugins import plugin_manager
+from httpie.sessions import get_httpie_session
 from httpie.utils import repr_dict
 
 
@@ -24,12 +26,90 @@ try:
 except (ImportError, AttributeError):
     pass
 
+
 FORM_CONTENT_TYPE = 'application/x-www-form-urlencoded; charset=utf-8'
 JSON_CONTENT_TYPE = 'application/json'
 JSON_ACCEPT = f'{JSON_CONTENT_TYPE}, */*'
 DEFAULT_UA = f'HTTPie/{__version__}'
 
 
+def collect_messages(
+    args: argparse.Namespace,
+    config_dir: Path
+) -> Iterable[Union[requests.PreparedRequest, requests.Response]]:
+    httpie_session = None
+    httpie_session_headers = None
+    if args.session or args.session_read_only:
+        httpie_session = get_httpie_session(
+            config_dir=config_dir,
+            session_name=args.session or args.session_read_only,
+            host=args.headers.get('Host'),
+            url=args.url,
+        )
+        httpie_session_headers = httpie_session.headers
+
+    request_kwargs = make_request_kwargs(
+        args=args,
+        base_headers=httpie_session_headers,
+    )
+    send_kwargs = make_send_kwargs(args)
+    send_kwargs_mergeable_from_env = make_send_kwargs_mergeable_from_env(args)
+    requests_session = build_requests_session(
+        ssl_version=args.ssl_version,
+        compress_arg=args.compress,
+    )
+
+    if httpie_session:
+        httpie_session.update_headers(request_kwargs['headers'])
+        requests_session.cookies = httpie_session.cookies
+        if args.auth_plugin:
+            # Save auth from CLI to HTTPie session.
+            httpie_session.auth = {
+                'type': args.auth_plugin.auth_type,
+                'raw_auth': args.auth_plugin.raw_auth,
+            }
+        elif httpie_session.auth:
+            # Apply auth from HTTPie session
+            request_kwargs['auth'] = httpie_session.auth
+
+    if args.debug:
+        # TODO: reflect the split between request and send kwargs.
+        dump_request(request_kwargs)
+
+    request = requests.Request(**request_kwargs)
+    prepared_request = requests_session.prepare_request(request)
+    response_count = 0
+    while prepared_request:
+        yield prepared_request
+        if not args.offline:
+            send_kwargs_merged = requests_session.merge_environment_settings(
+                url=prepared_request.url,
+                **send_kwargs_mergeable_from_env,
+            )
+            with max_headers(args.max_headers):
+                response = requests_session.send(
+                    request=prepared_request,
+                    **send_kwargs_merged,
+                    **send_kwargs,
+                )
+            response_count += 1
+            if response.next:
+                if args.max_redirects and response_count == args.max_redirects:
+                    raise requests.TooManyRedirects
+                if args.follow:
+                    prepared_request = response.next
+                    if args.all:
+                        yield response
+                    continue
+            yield response
+        break
+
+    if httpie_session:
+        if httpie_session.is_new() or not args.session_read_only:
+            httpie_session.cookies = requests_session.cookies
+            httpie_session.save()
+
+
 # noinspection PyProtectedMember
 @contextmanager
 def max_headers(limit):
@@ -83,14 +163,14 @@ class HTTPieHTTPAdapter(HTTPAdapter):
 
 
 def build_requests_session(
-    ssl_version: str,
     compress_arg: int,
+    ssl_version: str = None,
 ) -> requests.Session:
     requests_session = requests.Session()
 
     # Install our adapter.
     adapter = HTTPieHTTPAdapter(
-        ssl_version=ssl_version,
+        ssl_version=SSL_VERSION_ARG_MAPPING[ssl_version] if ssl_version else None,
         compression_enabled=compress_arg > 0,
         compress_always=compress_arg > 1,
     )
@@ -108,40 +188,6 @@ def build_requests_session(
     return requests_session
 
 
-def get_response(
-    args: argparse.Namespace,
-    config_dir: Path
-) -> requests.Response:
-    """Send the request and return a `request.Response`."""
-
-    ssl_version = None
-    if args.ssl_version:
-        ssl_version = SSL_VERSION_ARG_MAPPING[args.ssl_version]
-
-    requests_session = build_requests_session(
-        ssl_version=ssl_version,
-        compress_arg=args.compress
-    )
-    requests_session.max_redirects = args.max_redirects
-
-    with max_headers(args.max_headers):
-        if not args.session and not args.session_read_only:
-            requests_kwargs = make_requests_kwargs(args)
-            if args.debug:
-                dump_request(requests_kwargs)
-            response = requests_session.request(**requests_kwargs)
-        else:
-            response = sessions.get_response(
-                requests_session=requests_session,
-                args=args,
-                config_dir=config_dir,
-                session_name=args.session or args.session_read_only,
-                read_only=bool(args.session_read_only),
-            )
-
-    return response
-
-
 def dump_request(kwargs: dict):
     sys.stderr.write(
         f'\n>>> requests.request(**{repr_dict(kwargs)})\n\n')
@@ -181,12 +227,40 @@ def make_default_headers(args: argparse.Namespace) -> RequestHeadersDict:
     return default_headers
 
 
-def make_requests_kwargs(
+def make_send_kwargs(args: argparse.Namespace) -> dict:
+    kwargs = {
+        'timeout': args.timeout or None,
+        'allow_redirects': False,
+    }
+    return kwargs
+
+
+def make_send_kwargs_mergeable_from_env(args: argparse.Namespace) -> dict:
+    cert = None
+    if args.cert:
+        cert = args.cert
+        if args.cert_key:
+            cert = cert, args.cert_key
+    kwargs = {
+        'proxies': {p.key: p.value for p in args.proxy},
+        'stream': True,
+        'verify': {
+            'yes': True,
+            'true': True,
+            'no': False,
+            'false': False,
+        }.get(args.verify.lower(), args.verify),
+        'cert': cert,
+    }
+    return kwargs
+
+
+def make_request_kwargs(
     args: argparse.Namespace,
     base_headers: RequestHeadersDict = None
 ) -> dict:
     """
-    Translate our `args` into `requests.request` keyword arguments.
+    Translate our `args` into `requests.Request` keyword arguments.
 
     """
     # Serialize JSON data, if needed.
@@ -207,31 +281,14 @@ def make_requests_kwargs(
     headers.update(args.headers)
     headers = finalize_headers(headers)
 
-    cert = None
-    if args.cert:
-        cert = args.cert
-        if args.cert_key:
-            cert = cert, args.cert_key
-
     kwargs = {
-        'stream': True,
         'method': args.method.lower(),
         'url': args.url,
         'headers': headers,
         'data': data,
-        'verify': {
-            'yes': True,
-            'true': True,
-            'no': False,
-            'false': False,
-        }.get(args.verify.lower(), args.verify),
-        'cert': cert,
-        'timeout': args.timeout or None,
         'auth': args.auth,
-        'proxies': {p.key: p.value for p in args.proxy},
-        'files': args.files,
-        'allow_redirects': args.follow,
         'params': args.params,
+        'files': args.files,
     }
 
     return kwargs