From 474b1c1b679caf2fa0d07b51cde3ed0d6ab5fbf9 Mon Sep 17 00:00:00 2001
From: David Christenson <mail@noctem.xyz>
Date: Thu, 2 Mar 2017 21:37:19 -0700
Subject: [PATCH] Use ujson if available

Over 3x faster at encoding requests and 6x faster decoding responses in
my initial tests.
---
 aiopogo/auth_ptc.py    | 17 ++++++++++++-----
 aiopogo/hash_server.py | 25 +++++++++++++++++--------
 setup.py               |  1 +
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/aiopogo/auth_ptc.py b/aiopogo/auth_ptc.py
index a00595e5..e6387632 100755
--- a/aiopogo/auth_ptc.py
+++ b/aiopogo/auth_ptc.py
@@ -1,5 +1,4 @@
 from urllib.parse import parse_qs, urlsplit
-from json import JSONDecodeError
 from asyncio import get_event_loop, TimeoutError
 
 from aiohttp import TCPConnector, ClientSession, ClientError, DisconnectedError, HttpProcessingError
@@ -9,6 +8,14 @@
 from .utilities import get_time
 from .exceptions import ActivationRequiredException, AuthConnectionException, AuthException, AuthTimeoutException, InvalidCredentialsException, ProxyException, TimeoutException
 
+try:
+    import ujson as json
+
+    jexc = ValueError
+except ImportError:
+    import json
+
+    jexc = jexc = (json.JSONDecodeError, ValueError)
 
 class AuthPtc(Auth):
 
@@ -64,7 +71,7 @@ async def user_login(self, username=None, password=None, retry=True):
             now = get_time()
             async with self._session.get(self.PTC_LOGIN_URL, timeout=self.timeout, proxy=self.proxy) as resp:
                 resp.raise_for_status()
-                data = await resp.json()
+                data = await resp.json(loads=json.loads)
 
             try:
                 data['_eventId'] = 'submit'
@@ -81,8 +88,8 @@ async def user_login(self, username=None, password=None, retry=True):
                     self._access_token = resp.cookies['CASTGC'].value
                 except KeyError:
                     try:
-                        j = await resp.json()
-                    except JSONDecodeError as e:
+                        j = await resp.json(loads=json.loads)
+                    except jexc as e:
                         raise AuthException('Unable to decode second response.') from e
                     try:
                         if j.get('error_code') == 'users.login.activation_required':
@@ -105,7 +112,7 @@ async def user_login(self, username=None, password=None, retry=True):
             raise AuthTimeoutException('user_login timeout.') from e
         except ProxyException as e:
             raise ProxyException('Proxy connection error during user_login.') from e
-        except JSONDecodeError as e:
+        except jexc as e:
             raise AuthException('Unable to parse user_login response.') from e
         except (ClientError, DisconnectedError) as e:
             err = e.__cause__ or e
diff --git a/aiopogo/hash_server.py b/aiopogo/hash_server.py
index a3c77eaa..9a02998c 100644
--- a/aiopogo/hash_server.py
+++ b/aiopogo/hash_server.py
@@ -1,5 +1,3 @@
-import json
-
 from ctypes import c_int32, c_int64
 from base64 import b64encode
 from asyncio import get_event_loop, TimeoutError
@@ -7,9 +5,20 @@
 from aiohttp import ClientSession, ClientError, DisconnectedError, HttpProcessingError
 
 from .exceptions import ExpiredHashKeyException, HashingOfflineException, HashingQuotaExceededException, HashingTimeoutException, MalformedHashResponseException, TempHashingBanException, TimeoutException, UnexpectedHashResponseException
-from .utilities import JSONByteEncoder
 from .connector import TimedConnector
 
+try:
+    import ujson as json
+
+    jargs = {'escape_forward_slashes': False}
+    jexc = ValueError
+except ImportError:
+    import json
+    from .utilities import JSONByteEncoder
+
+    jargs = {'cls': JSONByteEncoder}
+    jexc = (json.JSONDecodeError, ValueError)
+
 
 class HashServer:
     endpoint = "https://pokehash.buddyauth.com/api/v127_3/hash"
@@ -21,7 +30,7 @@ def __init__(self, auth_token):
         self.auth_token = auth_token
         self.activate_session()
 
-    async def hash(self, timestamp, latitude, longitude, altitude, authticket, sessiondata, requests):
+    async def hash(self, timestamp, latitude, longitude, accuracy, authticket, sessiondata, requests):
         self.location_hash = None
         self.location_auth_hash = None
         headers = {'X-AuthToken': self.auth_token}
@@ -30,12 +39,12 @@ async def hash(self, timestamp, latitude, longitude, altitude, authticket, sessi
             'Timestamp': timestamp,
             'Latitude': latitude,
             'Longitude': longitude,
-            'Altitude': altitude,
+            'Altitude': accuracy,
             'AuthTicket': b64encode(authticket),
             'SessionData': b64encode(sessiondata),
             'Requests': tuple(b64encode(x.SerializeToString()) for x in requests)
         }
-        payload = json.dumps(payload, cls=JSONByteEncoder)
+        payload = json.dumps(payload, **jargs)
 
         # request hashes from hashing server
         try:
@@ -65,8 +74,8 @@ async def hash(self, timestamp, latitude, longitude, altitude, authticket, sessi
                     pass
 
                 try:
-                    response = await resp.json()
-                except (json.JSONDecodeError, ValueError) as e:
+                    response = await resp.json(encoding='ascii', loads=json.loads)
+                except jexc as e:
                     raise MalformedHashResponseException('Unable to parse JSON from hash server.') from e
         except (TimeoutException, TimeoutError) as e:
             raise HashingTimeoutException('Hashing request timed out.') from e
diff --git a/setup.py b/setup.py
index bbb2a81f..ea76b800 100755
--- a/setup.py
+++ b/setup.py
@@ -16,6 +16,7 @@
           'aiohttp==1.3.*',
           'pycrypt>=0.1.1',
           'pogeo>=0.2.0'],
+      extras_require={'ujson': ['ujson']},
       package_data={'aiopogo': ['lib/*.so', 'lib/*.dylib', 'lib/*.dll']},
       classifiers=[
           'Development Status :: 5 - Production/Stable',