Commit 6641c0c4 authored by Taddeüs Kroes's avatar Taddeüs Kroes

Code cleanup

parent bd6d116d
.PHONY: test clean .PHONY: check clean
test: check:
@python test.py @python test.py
clean: clean:
rm `find -name \*.pyc` find -name \*.pyc -delete
import struct import struct
from frame import ControlFrame, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG from frame import ControlFrame, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG, \
OPCODE_CONTINUATION
from message import create_message from message import create_message
from exceptions import SocketClosed, PingError from exceptions import SocketClosed, PingError
...@@ -43,8 +44,8 @@ class Connection(object): ...@@ -43,8 +44,8 @@ class Connection(object):
""" """
Receive a message. A message may consist of multiple (ordered) data Receive a message. A message may consist of multiple (ordered) data
frames. A control frame may be delivered at any time, also when frames. A control frame may be delivered at any time, also when
expecting the next data frame of a fragmented message. These control expecting the next continuation frame of a fragmented message. These
frames are handled immediately bu handle_control_frame(). control frames are handled immediately by handle_control_frame().
""" """
fragments = [] fragments = []
...@@ -57,10 +58,13 @@ class Connection(object): ...@@ -57,10 +58,13 @@ class Connection(object):
# No more receiving data after a close message # No more receiving data after a close message
if frame.opcode == OPCODE_CLOSE: if frame.opcode == OPCODE_CLOSE:
break break
elif len(fragments) and frame.opcode != OPCODE_CONTINUATION:
raise ValueError('expected continuation/control frame, got %s '
'instead' % frame)
else: else:
fragments.append(frame) fragments.append(frame)
payload = ''.join([f.payload for f in fragments]) payload = ''.join(f.payload for f in fragments)
return create_message(fragments[0].opcode, payload) return create_message(fragments[0].opcode, payload)
def handle_control_frame(self, frame): def handle_control_frame(self, frame):
......
...@@ -127,7 +127,7 @@ class Frame(object): ...@@ -127,7 +127,7 @@ class Frame(object):
""" """
frames = [] frames = []
for start in range(0, len(self.payload), fragment_size): for start in xrange(0, len(self.payload), fragment_size):
payload = self.payload[start:start + fragment_size] payload = self.payload[start:start + fragment_size]
frames.append(Frame(OPCODE_CONTINUATION, payload, mask=mask, frames.append(Frame(OPCODE_CONTINUATION, payload, mask=mask,
final=False)) final=False))
...@@ -149,7 +149,7 @@ class Frame(object): ...@@ -149,7 +149,7 @@ class Frame(object):
class ControlFrame(Frame): class ControlFrame(Frame):
""" """
A Control frame is a frame with an opcode OPCODE_CLOSE, OPCODE_PING or A control frame is a frame with an opcode OPCODE_CLOSE, OPCODE_PING or
OPCODE_PONG. These frames must be handled as defined by RFC 6455, and OPCODE_PONG. These frames must be handled as defined by RFC 6455, and
""" """
def fragment(self, fragment_size, mask=False): def fragment(self, fragment_size, mask=False):
......
...@@ -33,21 +33,21 @@ class BinaryMessage(Message): ...@@ -33,21 +33,21 @@ class BinaryMessage(Message):
class JSONMessage(TextMessage): class JSONMessage(TextMessage):
def __init__(self, dictionary, **kwargs): def __init__(self, data, **kwargs):
self.data = {} self.data = {}
self.data.extend(dictionary) self.data.update(data, **kwargs)
self.data.extend(kwargs)
super(JSONMessage, self).__init__(json.dumps(self.data)) super(JSONMessage, self).__init__(json.dumps(self.data))
@classmethod
OPCODE_CLASS_MAP = { def decode(cls, payload):
OPCODE_TEXT: TextMessage, return cls(json.loads(payload))
OPCODE_BINARY: BinaryMessage,
}
def create_message(opcode, payload): def create_message(opcode, payload):
if opcode in OPCODE_CLASS_MAP: if opcode == OPCODE_TEXT:
return OPCODE_CLASS_MAP[opcode](payload) return TextMessage(payload)
if opcode == OPCODE_BINARY:
return BinaryMessage(payload)
return Message(opcode, payload) return Message(opcode, payload)
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment