1 files changed, 23 insertions, 19 deletions
diff --git a/prolog/engine.py b/prolog/engine.py
index e0f8bf7..b023c9e 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -21,6 +21,7 @@ import http.client
 import json
 from operator import itemgetter
 import re
+import socket  # for socket.timeout exception
 import time
 import urllib3
 
@@ -59,26 +60,29 @@ def send(engine, event, timeout=10):
 
 # Return the main reply and pull potential output replies.
 def request(method, path, params=None, timeout=10):
-    if method == 'GET':
-        response = conn.request_encode_url(method, path, fields=params, timeout=timeout)
-    else:
-        response = conn.urlopen(method, path, body=params, timeout=timeout)
-
-    messages = []
-    while True:
-        if response.status != http.client.OK:
-            raise Exception('server returned {}'.format(response.status))
-
-        reply = json.loads(response.data.decode('utf-8'))
-        if isinstance(reply, dict) and reply.get('event') == 'output':
-            messages.append(_get_message(reply))
-
-            # Pull the next output. These requests should return instantly
-            # as no additional processing needs to be done in the pengine.
-            params = {'id': reply['id'], 'format': 'json-html'}
-            response = conn.request_encode_url('GET', '/pengine/pull_response', fields=params)
+    try:
+        if method == 'GET':
+            response = conn.request_encode_url(method, path, fields=params, timeout=timeout)
         else:
-            return reply, messages
+            response = conn.urlopen(method, path, body=params, timeout=timeout)
+
+        messages = []
+        while True:
+            if response.status != http.client.OK:
+                raise Exception('server returned {}'.format(response.status))
+
+            reply = json.loads(response.data.decode('utf-8'))
+            if isinstance(reply, dict) and reply.get('event') == 'output':
+                messages.append(_get_message(reply))
+
+                # Pull the next output. These requests should return instantly
+                # as no additional processing needs to be done in the pengine.
+                params = {'id': reply['id'], 'format': 'json-html'}
+                response = conn.request_encode_url('GET', '/pengine/pull_response', fields=params)
+            else:
+                return reply, messages
+    except urllib3.exceptions.ReadTimeoutError:
+        raise socket.timeout
 
 # Strip boilerplate from Prolog messages … ugly.
 def _get_message(reply):