Add Logger (#184)

* add logger

* pretty logging

* lint

Author: clickingbuttons

polygon/logging.py

@@ -0,0 +1,12 @@
+import logging
+import sys
+
+
+def get_logger(name: str) -> logging.Logger:
+    logger = logging.getLogger(name)
+    handler = logging.StreamHandler(sys.stdout)
+    formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    return logger

polygon/rest/base.py

@@ -6,7 +6,10 @@
 from typing import Optional, Any, Dict
 from datetime import datetime
 import pkg_resources  # part of setuptools
+from ..logging import get_logger
+import logging
 
+logger = get_logger("RESTClient")
 version = "unknown"
 try:
     version = pkg_resources.require("polygon-api-client")[0].version
@@ -49,7 +52,8 @@ def __init__(
         )
         self.timeout = urllib3.Timeout(connect=connect_timeout, read=read_timeout)
         self.retries = retries
-        self.verbose = verbose
+        if verbose:
+            logger.setLevel(logging.DEBUG)
 
     def _decode(self, resp):
         return json.loads(resp.data.decode("utf-8"))
@@ -65,8 +69,7 @@ def _get(
         if params is None:
             params = {}
         params = {str(k): str(v) for k, v in params.items() if v is not None}
-        if self.verbose:
-            print("_get", path, params)
+        logger.debug("_get %s params %s", path, params)
         resp = self.client.request(
             "GET",
             self.BASE + path,
@@ -144,7 +147,6 @@ def _paginate_iter(
         self,
         path: str,
         params: dict,
-        raw: bool,
         deserializer,
         result_key: str = "results",
     ):
@@ -183,5 +185,4 @@ def _paginate(
             params=params,
             deserializer=deserializer,
             result_key=result_key,
-            raw=True,
         )

polygon/websocket/__init__.py

@@ -1,15 +1,19 @@
 import os
 from enum import Enum
 from typing import Optional, Union, List, Set, Callable, Awaitable
+import logging
 import json
 import asyncio
 import ssl
 import certifi
 from .models import *
 from websockets.client import connect, WebSocketClientProtocol
 from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError
+from ..logging import get_logger
+import logging
 
 env_key = "POLYGON_API_KEY"
+logger = get_logger("WebSocketClient")
 
 
 class AuthError(Exception):
@@ -35,7 +39,7 @@ def __init__(
         :param api_key: Your API keYour API key.
         :param feed: The feed to subscribe to.
         :param raw: Whether to pass raw Union[str, bytes] to user callback.
-        :param verbose: Whether to print client and server status messages.
+        :param verbose: Whether to log client and server status messages.
         :param subscriptions: List of subscription parameters.
         :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop.
         :return: A client.
@@ -48,7 +52,8 @@ def __init__(
         self.feed = feed
         self.market = market
         self.raw = raw
-        self.verbose = verbose
+        if verbose:
+            logger.setLevel(logging.DEBUG)
         self.websocket_cfg = kwargs
         if isinstance(feed, Enum):
             feed = feed.value
@@ -82,8 +87,7 @@ async def connect(
         :raises AuthError: If invalid API key is supplied.
         """
         reconnects = 0
-        if self.verbose:
-            print("connect:", self.url)
+        logger.debug("connect: %s", self.url)
         # darwin needs some extra <3
         ssl_context = None
         if self.url.startswith("wss://"):
@@ -96,21 +100,19 @@ async def connect(
             self.websocket = s
             try:
                 msg = await s.recv()
-                if self.verbose:
-                    print("connected:", msg)
-                if self.verbose:
-                    print("authing:")
+                logger.debug("connected: %s", msg)
+                logger.debug("authing...")
                 await s.send(json.dumps({"action": "auth", "params": self.api_key}))
                 auth_msg = await s.recv()
                 auth_msg_parsed = json.loads(auth_msg)
-                if self.verbose:
-                    print("authed:", auth_msg)
+                logger.debug("authed: %s", auth_msg)
                 if auth_msg_parsed[0]["status"] == "auth_failed":
                     raise AuthError(auth_msg_parsed[0]["message"])
                 while True:
                     if self.schedule_resub:
-                        if self.verbose:
-                            print("reconciling:", self.subs, self.scheduled_subs)
+                        logger.debug(
+                            "reconciling: %s %s", self.subs, self.scheduled_subs
+                        )
                         new_subs = self.scheduled_subs.difference(self.subs)
                         await self._subscribe(new_subs)
                         old_subs = self.subs.difference(self.scheduled_subs)
@@ -126,21 +128,18 @@ async def connect(
                     msgJson = json.loads(cmsg)  # type: ignore
                     for m in msgJson:
                         if m["ev"] == "status":
-                            if self.verbose:
-                                print("status:", m["message"])
+                            logger.debug("status: %s", m["message"])
                             continue
                     if not self.raw:
-                        cmsg = parse(msgJson)
+                        cmsg = parse(msgJson, logger)
 
                     if len(cmsg) > 0:
                         await processor(cmsg)  # type: ignore
             except ConnectionClosedOK:
-                if self.verbose:
-                    print("connection closed (OK)")
+                logger.debug("connection closed (OK)")
                 return
             except ConnectionClosedError:
-                if self.verbose:
-                    print("connection closed (error)")
+                logger.debug("connection closed (ERR)")
                 reconnects += 1
                 if self.max_reconnects is not None and reconnects > self.max_reconnects:
                     return
@@ -172,24 +171,22 @@ async def _subscribe(self, topics: Union[List[str], Set[str]]):
         if self.websocket is None or len(topics) == 0:
             return
         subs = ",".join(topics)
-        if self.verbose:
-            print("subbing:", subs)
+        logger.debug("subbing: %s", subs)
         await self.websocket.send(json.dumps({"action": "subscribe", "params": subs}))
 
     async def _unsubscribe(self, topics: Union[List[str], Set[str]]):
         if self.websocket is None or len(topics) == 0:
             return
         subs = ",".join(topics)
-        if self.verbose:
-            print("unsubbing:", subs)
+        logger.debug("unsubbing: %s", subs)
         await self.websocket.send(json.dumps({"action": "unsubscribe", "params": subs}))
 
     @staticmethod
     def _parse_subscription(s: str):
         s = s.strip()
         split = s.split(".")
         if len(split) != 2:
-            print("invalid subscription:", s)
+            logger.warning("invalid subscription:", s)
             return [None, None]
 
         return split
@@ -204,8 +201,7 @@ def subscribe(self, *subscriptions: str):
             topic, sym = self._parse_subscription(s)
             if topic == None:
                 continue
-            if self.verbose:
-                print("add:", s)
+            logger.debug("sub desired: %s", s)
             self.scheduled_subs.add(s)
             # If user subs to X.*, remove other X.\w+
             if sym == "*":
@@ -225,8 +221,7 @@ def unsubscribe(self, *subscriptions: str):
             topic, sym = self._parse_subscription(s)
             if topic == None:
                 continue
-            if self.verbose:
-                print("discard:", s)
+            logger.debug("sub undesired: %s", s)
             self.scheduled_subs.discard(s)
 
             # If user unsubs to X.*, remove other X.\w+
@@ -248,11 +243,10 @@ async def close(self):
         """
         Close the websocket connection.
         """
-        if self.verbose:
-            print("closing:")
+        logger.debug("closing")
 
         if self.websocket:
             await self.websocket.close()
             self.websocket = None
         else:
-            print("no websocket open to close")
+            logger.warning("no websocket open to close")

polygon/websocket/models/__init__.py

@@ -1,6 +1,7 @@
 from typing import Dict, Any, List
 from .common import *
 from .models import *
+import logging
 
 
 def parse_single(data: Dict[str, Any]):
@@ -30,13 +31,13 @@ def parse_single(data: Dict[str, Any]):
     return None
 
 
-def parse(msg: List[Dict[str, Any]]) -> List[WebSocketMessage]:
+def parse(msg: List[Dict[str, Any]], logger: logging.Logger) -> List[WebSocketMessage]:
     res = []
     for m in msg:
         parsed = parse_single(m)
         if parsed is None:
             if m["ev"] != "status":
-                print("could not parse message", m)
+                logger.warning("could not parse message %s", m)
         else:
             res.append(parsed)
     return res

test_websocket/base_ws.py

@@ -22,4 +22,7 @@ def expectResponse(self, msg):
     async def asyncSetUp(self):
         self.maxDiff = None
         loop = asyncio.get_event_loop()
-        loop.create_task(run_mock_server())
+        self.task = loop.create_task(run_mock_server())
+
+    async def asyncTearDown(self):
+        self.task.cancel()