Taddeus Kroes 12 лет назад
Родитель
Сommit
6641c0c452
4 измененных файлов с 23 добавлено и 19 удалено
  1. 3 3
      Makefile
  2. 8 4
      connection.py
  3. 2 2
      frame.py
  4. 10 10
      message.py

+ 3 - 3
Makefile

@@ -1,7 +1,7 @@
-.PHONY: test clean
+.PHONY: check clean
 
-test:
+check:
 	@python test.py
 
 clean:
-	rm `find -name \*.pyc`
+	find -name \*.pyc -delete

+ 8 - 4
connection.py

@@ -1,6 +1,7 @@
 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 exceptions import SocketClosed, PingError
 
@@ -43,8 +44,8 @@ class Connection(object):
         """
         Receive a message. A message may consist of multiple (ordered) data
         frames. A control frame may be delivered at any time, also when
-        expecting the next data frame of a fragmented message. These control
-        frames are handled immediately bu handle_control_frame().
+        expecting the next continuation frame of a fragmented message. These
+        control frames are handled immediately by handle_control_frame().
         """
         fragments = []
 
@@ -57,10 +58,13 @@ class Connection(object):
                 # No more receiving data after a close message
                 if frame.opcode == OPCODE_CLOSE:
                     break
+            elif len(fragments) and frame.opcode != OPCODE_CONTINUATION:
+                raise ValueError('expected continuation/control frame, got %s '
+                                 'instead' % frame)
             else:
                 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)
 
     def handle_control_frame(self, frame):

+ 2 - 2
frame.py

@@ -127,7 +127,7 @@ class Frame(object):
         """
         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]
             frames.append(Frame(OPCODE_CONTINUATION, payload, mask=mask,
                                 final=False))
@@ -149,7 +149,7 @@ class Frame(object):
 
 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
     """
     def fragment(self, fragment_size, mask=False):

+ 10 - 10
message.py

@@ -33,21 +33,21 @@ class BinaryMessage(Message):
 
 
 class JSONMessage(TextMessage):
-    def __init__(self, dictionary, **kwargs):
+    def __init__(self, data, **kwargs):
         self.data = {}
-        self.data.extend(dictionary)
-        self.data.extend(kwargs)
+        self.data.update(data, **kwargs)
         super(JSONMessage, self).__init__(json.dumps(self.data))
 
-
-OPCODE_CLASS_MAP = {
-    OPCODE_TEXT: TextMessage,
-    OPCODE_BINARY: BinaryMessage,
-}
+    @classmethod
+    def decode(cls, payload):
+        return cls(json.loads(payload))
 
 
 def create_message(opcode, payload):
-    if opcode in OPCODE_CLASS_MAP:
-        return OPCODE_CLASS_MAP[opcode](payload)
+    if opcode == OPCODE_TEXT:
+        return TextMessage(payload)
+
+    if opcode == OPCODE_BINARY:
+        return BinaryMessage(payload)
 
     return Message(opcode, payload)