From 8ff2a144b436ff8425363f5d6e83ce885cbf6816 Mon Sep 17 00:00:00 2001
From: Alex Dehnert <adehnert@mit.edu>
Date: Tue, 18 Jul 2017 03:36:07 -0400
Subject: [PATCH] api: Use requests.Session

Using requests.Session allows the requests library to reuse HTTP connections,
which is potentially helpful for performance. This fixes python-zulip-api #3.
---
 zulip/zulip/__init__.py | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py
index 18bd3260c..a78bd475b 100644
--- a/zulip/zulip/__init__.py
+++ b/zulip/zulip/__init__.py
@@ -285,6 +285,30 @@ def __init__(self, email=None, api_key=None, config_file=None,
         self.client_cert = client_cert
         self.client_cert_key = client_cert_key
 
+        self.session = None  # type: Union[None, requests.Session]
+
+    def ensure_session(self):
+        # type: () -> None
+
+        # Check if the session has been created already, and return
+        # immediately if so.
+        if self.session:
+            return
+
+        # Build a client cert object for requests
+        if self.client_cert_key is not None:
+            client_cert = (self.client_cert, self.client_cert_key)  # type: Union[str, Tuple[str, str]]
+        else:
+            client_cert = self.client_cert
+
+        # Actually construct the session
+        session = requests.Session()
+        session.auth = requests.auth.HTTPBasicAuth(self.email, self.api_key)  # type: ignore # https://github.com/python/typeshed/pull/1504
+        session.verify = self.tls_verification  # type: ignore # https://github.com/python/typeshed/pull/1504
+        session.cert = client_cert
+        session.headers = {"User-agent": self.get_user_agent()}
+        self.session = session
+
     def get_user_agent(self):
         # type: () -> str
         vendor = ''
@@ -327,6 +351,8 @@ def do_api_query(self, orig_request, url, method="POST", longpolling=False, file
         for f in files:
             req_files.append((f.name, f))
 
+        self.ensure_session()
+
         query_state = {
             'had_error_retry': False,
             'request': request,
@@ -370,21 +396,10 @@ def end_error_retry(succeeded):
                 if files:
                     kwargs['files'] = req_files
 
-                # Build a client cert object for requests
-                if self.client_cert_key is not None:
-                    client_cert = (self.client_cert, self.client_cert_key)  # type: Union[str, Tuple[str, str]]
-                else:
-                    client_cert = self.client_cert
-
-                res = requests.request(
+                res = self.session.request(
                     method,
                     urllib.parse.urljoin(self.base_url, url),
-                    auth=requests.auth.HTTPBasicAuth(self.email,
-                                                     self.api_key),
-                    verify=self.tls_verification,
-                    cert=client_cert,
                     timeout=90,
-                    headers={"User-agent": self.get_user_agent()},
                     **kwargs)
 
                 # On 50x errors, try again after a short sleep