Add type annotations (PEP 484) (#25)

This commit is contained in:
Sergey Petrov
2019-07-29 01:26:15 +03:00
committed by GitHub
parent be7438a653
commit fde28abb04
9 changed files with 151 additions and 76 deletions
+35 -26
View File
@@ -4,13 +4,15 @@ import logging
import time
import weakref
from enum import Enum
from json import JSONEncoder
from threading import Thread
from typing import Deque, Dict, Iterable, Optional, Tuple, Type, Union
from .credentials import CertificateCredentials, Credentials
from .errors import ConnectionFailed, exception_class_for_reason
# We don't generally need to know about the Credentials subclasses except to
# keep the old API, where APNsClient took a cert_file
from .credentials import CertificateCredentials
from .payload import Payload
class NotificationPriority(Enum):
@@ -35,10 +37,14 @@ class APNsClient(object):
DEFAULT_PORT = 443
ALTERNATIVE_PORT = 2197
def __init__(self, credentials, use_sandbox=False, use_alternative_port=False, proto=None, json_encoder=None,
password=None, proxy_host=None, proxy_port=None, heartbeat_period=None):
if credentials is None or isinstance(credentials, str):
self.__credentials = CertificateCredentials(credentials, password)
def __init__(self,
credentials: Union[Credentials, str],
use_sandbox: bool = False, use_alternative_port: bool = False, proto: Optional[str] = None,
json_encoder: Optional[Type[JSONEncoder]] = None, password: Optional[str] = None,
proxy_host: Optional[str] = None, proxy_port: Optional[int] = None,
heartbeat_period: Optional[float] = None) -> None:
if isinstance(credentials, str):
self.__credentials = CertificateCredentials(credentials, password) # type: Credentials
else:
self.__credentials = credentials
self._init_connection(use_sandbox, use_alternative_port, proto, proxy_host, proxy_port)
@@ -47,18 +53,19 @@ class APNsClient(object):
self._start_heartbeat(heartbeat_period)
self.__json_encoder = json_encoder
self.__max_concurrent_streams = None
self.__max_concurrent_streams = 0
self.__previous_server_max_concurrent_streams = None
def _init_connection(self, use_sandbox, use_alternative_port, proto, proxy_host, proxy_port):
def _init_connection(self, use_sandbox: bool, use_alternative_port: bool, proto: Optional[str],
proxy_host: Optional[str], proxy_port: Optional[int]) -> None:
server = self.SANDBOX_SERVER if use_sandbox else self.LIVE_SERVER
port = self.ALTERNATIVE_PORT if use_alternative_port else self.DEFAULT_PORT
self._connection = self.__credentials.create_connection(server, port, proto, proxy_host, proxy_port)
def _start_heartbeat(self, heartbeat_period):
def _start_heartbeat(self, heartbeat_period: float) -> None:
conn_ref = weakref.ref(self._connection)
def watchdog():
def watchdog() -> None:
while True:
conn = conn_ref()
if conn is None:
@@ -71,8 +78,9 @@ class APNsClient(object):
thread.setDaemon(True)
thread.start()
def send_notification(self, token_hex, notification, topic=None, priority=NotificationPriority.Immediate,
expiration=None, collapse_id=None):
def send_notification(self, token_hex: str, notification: Payload, topic: Optional[str] = None,
priority: NotificationPriority = NotificationPriority.Immediate,
expiration: Optional[int] = None, collapse_id: Optional[str] = None) -> None:
stream_id = self.send_notification_async(token_hex, notification, topic, priority, expiration, collapse_id)
result = self.get_notification_result(stream_id)
if result != 'Success':
@@ -82,8 +90,9 @@ class APNsClient(object):
else:
raise exception_class_for_reason(result)
def send_notification_async(self, token_hex, notification, topic=None, priority=NotificationPriority.Immediate,
expiration=None, collapse_id=None):
def send_notification_async(self, token_hex: str, notification: Payload, topic: Optional[str] = None,
priority: NotificationPriority = NotificationPriority.Immediate,
expiration: Optional[int] = None, collapse_id: Optional[str] = None) -> int:
json_str = json.dumps(notification.dict(), cls=self.__json_encoder, ensure_ascii=False, separators=(',', ':'))
json_payload = json_str.encode('utf-8')
@@ -105,10 +114,10 @@ class APNsClient(object):
headers['apns-collapse-id'] = collapse_id
url = '/3/device/{}'.format(token_hex)
stream_id = self._connection.request('POST', url, json_payload, headers)
stream_id = self._connection.request('POST', url, json_payload, headers) # type: int
return stream_id
def get_notification_result(self, stream_id):
def get_notification_result(self, stream_id: int) -> Union[str, Tuple[str, str]]:
"""
Get result for specified stream
The function returns: 'Success' or 'failure reason' or ('Unregistered', timestamp)
@@ -118,14 +127,16 @@ class APNsClient(object):
return 'Success'
else:
raw_data = response.read().decode('utf-8')
data = json.loads(raw_data)
data = json.loads(raw_data) # type: Dict[str, str]
if response.status == 410:
return data['reason'], data['timestamp']
else:
return data['reason']
def send_notification_batch(self, notifications, topic=None, priority=NotificationPriority.Immediate,
expiration=None, collapse_id=None):
def send_notification_batch(self, notifications: Iterable[Notification], topic: Optional[str] = None,
priority: NotificationPriority = NotificationPriority.Immediate,
expiration: Optional[int] = None, collapse_id: Optional[str] = None
) -> Dict[str, Union[str, Tuple[str, str]]]:
"""
Send a notification to a list of tokens in batch. Instead of sending a synchronous request
for each token, send multiple requests concurrently. This is done on the same connection,
@@ -146,7 +157,7 @@ class APNsClient(object):
self.connect()
results = {}
open_streams = collections.deque()
open_streams = collections.deque() # type: Deque[RequestStream]
# Loop on the tokens, sending as many requests as possible concurrently to APNs.
# When reaching the maximum concurrent streams limit, wait for a response before sending
# another request.
@@ -154,7 +165,7 @@ class APNsClient(object):
# Update the max_concurrent_streams on every iteration since a SETTINGS frame can be
# sent by the server at any time.
self.update_max_concurrent_streams()
if self.should_send_notification(next_notification, open_streams):
if next_notification is not None and len(open_streams) < self.__max_concurrent_streams:
logger.info('Sending to token %s', next_notification.token)
stream_id = self.send_notification_async(next_notification.token, next_notification.payload, topic,
priority, expiration, collapse_id)
@@ -175,10 +186,7 @@ class APNsClient(object):
return results
def should_send_notification(self, notification, open_streams):
return notification is not None and len(open_streams) < self.__max_concurrent_streams
def update_max_concurrent_streams(self):
def update_max_concurrent_streams(self) -> None:
# Get the max_concurrent_streams setting returned by the server.
# The max_concurrent_streams value is saved in the H2Connection instance that must be
# accessed using a with statement in order to acquire a lock.
@@ -204,13 +212,14 @@ class APNsClient(object):
logger.info('APNs set max_concurrent_streams to %s', max_concurrent_streams)
self.__max_concurrent_streams = max_concurrent_streams
def connect(self):
def connect(self) -> None:
"""
Establish a connection to APNs. If already connected, the function does nothing. If the
connection fails, the function retries up to MAX_CONNECTION_RETRIES times.
"""
retries = 0
while retries < MAX_CONNECTION_RETRIES:
# noinspection PyBroadException
try:
self._connection.connect()
logger.info('Connected to APNs')
+23 -16
View File
@@ -1,8 +1,13 @@
import time
from typing import Optional, Tuple, TYPE_CHECKING
import jwt
from hyper import HTTP20Connection
from hyper.tls import init_context
from hyper import HTTP20Connection # type: ignore
from hyper.tls import init_context # type: ignore
if TYPE_CHECKING:
from hyper.ssl_compat import SSLContext # type: ignore
DEFAULT_TOKEN_LIFETIME = 3600
DEFAULT_TOKEN_ENCRYPTION_ALGORITHM = 'ES256'
@@ -10,22 +15,25 @@ DEFAULT_TOKEN_ENCRYPTION_ALGORITHM = 'ES256'
# Abstract Base class. This should not be instantiated directly.
class Credentials(object):
def __init__(self, ssl_context=None):
def __init__(self, ssl_context: 'Optional[SSLContext]' = None) -> None:
super().__init__()
self.__ssl_context = ssl_context
# Creates a connection with the credentials, if available or necessary.
def create_connection(self, server, port, proto, proxy_host=None, proxy_port=None):
def create_connection(self, server: str, port: int, proto: Optional[str], proxy_host: Optional[str] = None,
proxy_port: Optional[int] = None) -> HTTP20Connection:
# self.__ssl_context may be none, and that's fine.
return HTTP20Connection(server, port, ssl_context=self.__ssl_context, force_proto=proto or 'h2',
secure=True, proxy_host=proxy_host, proxy_port=proxy_port)
def get_authorization_header(self, topic):
def get_authorization_header(self, topic: Optional[str]) -> Optional[str]:
return None
# Credentials subclass for certificate authentication
class CertificateCredentials(Credentials):
def __init__(self, cert_file=None, password=None, cert_chain=None):
def __init__(self, cert_file: Optional[str] = None, password: Optional[str] = None,
cert_chain: Optional[str] = None) -> None:
ssl_context = init_context(cert=cert_file, cert_password=password)
if cert_chain:
ssl_context.load_cert_chain(cert_chain)
@@ -34,38 +42,37 @@ class CertificateCredentials(Credentials):
# Credentials subclass for JWT token based authentication
class TokenCredentials(Credentials):
def __init__(self, auth_key_path, auth_key_id, team_id,
encryption_algorithm=DEFAULT_TOKEN_ENCRYPTION_ALGORITHM,
token_lifetime=DEFAULT_TOKEN_LIFETIME):
def __init__(self, auth_key_path: str, auth_key_id: str, team_id: str,
encryption_algorithm: str = DEFAULT_TOKEN_ENCRYPTION_ALGORITHM,
token_lifetime: int = DEFAULT_TOKEN_LIFETIME) -> None:
self.__auth_key = self._get_signing_key(auth_key_path)
self.__auth_key_id = auth_key_id
self.__team_id = team_id
self.__encryption_algorithm = encryption_algorithm
self.__token_lifetime = token_lifetime
# Dictionary of {topic: (issue time, ascii decoded token)}
self.__jwt_token = None
self.__jwt_token = None # type: Optional[Tuple[float, str]]
# Use the default constructor because we don't have an SSL context
super(TokenCredentials, self).__init__()
def get_authorization_header(self, topic):
token = self._get_or_create_topic_token(topic)
def get_authorization_header(self, topic: Optional[str]) -> str:
token = self._get_or_create_topic_token()
return 'bearer %s' % token
@staticmethod
def _is_expired_token(issue_date):
def _is_expired_token(issue_date: float) -> bool:
return time.time() > issue_date + DEFAULT_TOKEN_LIFETIME
@staticmethod
def _get_signing_key(key_path):
def _get_signing_key(key_path: str) -> str:
secret = ''
if key_path:
with open(key_path) as f:
secret = f.read()
return secret
def _get_or_create_topic_token(self, topic):
def _get_or_create_topic_token(self) -> str:
# dict of topic to issue date and JWT token
token_pair = self.__jwt_token
if token_pair is None or self._is_expired_token(token_pair[0]):
+5 -2
View File
@@ -1,3 +1,6 @@
from typing import Type, Optional
class APNsException(Exception):
pass
@@ -128,7 +131,7 @@ class MethodNotAllowed(InternalException):
class Unregistered(APNsException):
"""The device token is inactive for the specified topic."""
def __init__(self, timestamp=None):
def __init__(self, timestamp: Optional[str] = None) -> None:
super(Unregistered, self).__init__()
self.timestamp = timestamp
@@ -164,7 +167,7 @@ class Shutdown(APNsException):
pass
def exception_class_for_reason(reason):
def exception_class_for_reason(reason: str) -> Type[APNsException]:
return {
'BadCollapseId': BadCollapseId,
'BadDeviceToken': BadDeviceToken,
+31 -11
View File
@@ -1,11 +1,21 @@
from typing import Any, Dict, List, Optional, Union, Iterable
MAX_PAYLOAD_SIZE = 4096
class PayloadAlert(object):
def __init__(self, title=None, title_localized_key=None, title_localized_args=None,
body=None, body_localized_key=None, body_localized_args=None,
action_localized_key=None, action=None,
launch_image=None):
def __init__(
self,
title: Optional[str] = None,
title_localized_key: Optional[str] = None,
title_localized_args: Optional[List[str]] = None,
body: Optional[str] = None,
body_localized_key: Optional[str] = None,
body_localized_args: Optional[List[str]] = None,
action_localized_key: Optional[str] = None,
action: Optional[str] = None,
launch_image: Optional[str] = None
) -> None:
self.title = title
self.title_localized_key = title_localized_key
self.title_localized_args = title_localized_args
@@ -16,8 +26,8 @@ class PayloadAlert(object):
self.action = action
self.launch_image = launch_image
def dict(self):
result = {}
def dict(self) -> Dict[str, Any]:
result = {} # type: Dict[str, Any]
if self.title:
result['title'] = self.title
@@ -45,9 +55,18 @@ class PayloadAlert(object):
class Payload(object):
def __init__(self, alert=None, badge=None, sound=None,
content_available=False, mutable_content=False,
category=None, url_args=None, custom=None, thread_id=None):
def __init__(
self,
alert: Union[PayloadAlert, str, None] = None,
badge: Optional[int] = None,
sound: Optional[str] = None,
category: Optional[str] = None,
url_args: Optional[Iterable[str]] = None,
custom: Optional[Dict[str, Any]] = None,
thread_id: Optional[str] = None,
content_available: bool = False,
mutable_content: bool = False,
) -> None:
self.alert = alert
self.badge = badge
self.sound = sound
@@ -58,10 +77,11 @@ class Payload(object):
self.mutable_content = mutable_content
self.thread_id = thread_id
def dict(self):
def dict(self) -> Dict[str, Any]:
result = {
'aps': {}
}
} # type: Dict[str, Any]
if self.alert is not None:
if isinstance(self.alert, PayloadAlert):
result['aps']['alert'] = self.alert.dict()