Factorize VariableHeader with packet id

hbmqtt/mqtt/packet.py

@@ -210,3 +210,19 @@ class MQTTPacket:
 
     def __repr__(self):
         return 'MQPacket(fixed={0!r}, variable={1!r}, payload={2!r})'.format(self.fixed_header, self.variable_header, self.payload)
+
+
+class PacketIdVariableHeader(MQTTVariableHeader):
+    def __init__(self, packet_id):
+        super().__init__()
+        self.packet_id = packet_id
+
+    def to_bytes(self):
+        out = b''
+        out += int_to_bytes(self.packet_id, 2)
+        return out
+
+    @classmethod
+    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
+        packet_id = yield from decode_packet_id(reader)
+        return cls(packet_id)

hbmqtt/mqtt/puback.py

@@ -1,32 +1,15 @@
 # Copyright (c) 2015 Nicolas JOUANIN
 #
 # See the file license.txt for copying permission.
-from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, PacketIdVariableHeader
 from hbmqtt.errors import HBMQTTException
-from hbmqtt.codecs import *
-
-
-class PubackVariableHeader(MQTTVariableHeader):
-    def __init__(self, packet_id):
-        super().__init__()
-        self.packet_id = packet_id
-
-    def to_bytes(self):
-        out = b''
-        out += int_to_bytes(self.packet_id, 2)
-        return out
-
-    @classmethod
-    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
-        packet_id = yield from decode_packet_id(reader)
-        return cls(packet_id)
 
 
 class PubackPacket(MQTTPacket):
-    VARIABLE_HEADER = PubackVariableHeader
+    VARIABLE_HEADER = PacketIdVariableHeader
     PAYLOAD = None
 
-    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubackVariableHeader=None, payload=None):
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PacketIdVariableHeader=None, payload=None):
         if fixed is None:
             header = MQTTFixedHeader(PacketType.PUBACK, 0x00)
         else:

hbmqtt/mqtt/pubcomp.py

@@ -1,32 +1,15 @@
 # Copyright (c) 2015 Nicolas JOUANIN
 #
 # See the file license.txt for copying permission.
-from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, PacketIdVariableHeader
 from hbmqtt.errors import HBMQTTException
-from hbmqtt.codecs import *
-
-
-class PubcompVariableHeader(MQTTVariableHeader):
-    def __init__(self, packet_id):
-        super().__init__()
-        self.packet_id = packet_id
-
-    def to_bytes(self):
-        out = b''
-        out += int_to_bytes(self.packet_id, 2)
-        return out
-
-    @classmethod
-    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
-        packet_id = yield from decode_packet_id(reader)
-        return cls(packet_id)
 
 
 class PubcompPacket(MQTTPacket):
-    VARIABLE_HEADER = PubcompVariableHeader
+    VARIABLE_HEADER = PacketIdVariableHeader
     PAYLOAD = None
 
-    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubcompVariableHeader=None, payload=None):
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PacketIdVariableHeader=None, payload=None):
         if fixed is None:
             header = MQTTFixedHeader(PacketType.PUBCOMP, 0x00)
         else:

hbmqtt/mqtt/pubrec.py

@@ -1,32 +1,15 @@
 # Copyright (c) 2015 Nicolas JOUANIN
 #
 # See the file license.txt for copying permission.
-from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, PacketIdVariableHeader
 from hbmqtt.errors import HBMQTTException
-from hbmqtt.codecs import *
-
-
-class PubrecVariableHeader(MQTTVariableHeader):
-    def __init__(self, packet_id):
-        super().__init__()
-        self.packet_id = packet_id
-
-    def to_bytes(self):
-        out = b''
-        out += int_to_bytes(self.packet_id, 2)
-        return out
-
-    @classmethod
-    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
-        packet_id = yield from decode_packet_id(reader)
-        return cls(packet_id)
 
 
 class PubrecPacket(MQTTPacket):
-    VARIABLE_HEADER = PubrecVariableHeader
+    VARIABLE_HEADER = PacketIdVariableHeader
     PAYLOAD = None
 
-    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubrecVariableHeader=None, payload=None):
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PacketIdVariableHeader=None, payload=None):
         if fixed is None:
             header = MQTTFixedHeader(PacketType.PUBREC, 0x00)
         else:

hbmqtt/mqtt/pubrel.py

@@ -1,32 +1,15 @@
 # Copyright (c) 2015 Nicolas JOUANIN
 #
 # See the file license.txt for copying permission.
-from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, PacketIdVariableHeader
 from hbmqtt.errors import HBMQTTException
-from hbmqtt.codecs import *
-
-
-class PubrelVariableHeader(MQTTVariableHeader):
-    def __init__(self, packet_id):
-        super().__init__()
-        self.packet_id = packet_id
-
-    def to_bytes(self):
-        out = b''
-        out += int_to_bytes(self.packet_id, 2)
-        return out
-
-    @classmethod
-    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
-        packet_id = yield from decode_packet_id(reader)
-        return cls(packet_id)
 
 
 class PubrelPacket(MQTTPacket):
-    VARIABLE_HEADER = PubrelVariableHeader
+    VARIABLE_HEADER = PacketIdVariableHeader
     PAYLOAD = None
 
-    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubrelVariableHeader=None, payload=None):
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PacketIdVariableHeader=None, payload=None):
         if fixed is None:
             header = MQTTFixedHeader(PacketType.PUBREL, 0x00)
         else:

tests/mqtt/test_puback.py

@@ -3,7 +3,7 @@
 # See the file license.txt for copying permission.
 import unittest
 
-from hbmqtt.mqtt.puback import PubackPacket, PubackVariableHeader
+from hbmqtt.mqtt.puback import PubackPacket, PacketIdVariableHeader
 from hbmqtt.codecs import *
 
 class PubackPacketTest(unittest.TestCase):
@@ -18,7 +18,7 @@ class PubackPacketTest(unittest.TestCase):
         self.assertEqual(message.variable_header.packet_id, 10)
 
     def test_to_bytes(self):
-        variable_header = PubackVariableHeader(10)
+        variable_header = PacketIdVariableHeader(10)
         publish = PubackPacket(variable_header=variable_header)
         out = publish.to_bytes()
         self.assertEqual(out, b'\x40\x02\x00\x0a')

tests/mqtt/test_pubcomp.py

@@ -3,7 +3,7 @@
 # See the file license.txt for copying permission.
 import unittest
 
-from hbmqtt.mqtt.pubcomp import PubcompPacket, PubcompVariableHeader
+from hbmqtt.mqtt.pubcomp import PubcompPacket, PacketIdVariableHeader
 from hbmqtt.codecs import *
 
 class PubcompPacketTest(unittest.TestCase):
@@ -18,7 +18,7 @@ class PubcompPacketTest(unittest.TestCase):
         self.assertEqual(message.variable_header.packet_id, 10)
 
     def test_to_bytes(self):
-        variable_header = PubcompVariableHeader(10)
+        variable_header = PacketIdVariableHeader(10)
         publish = PubcompPacket(variable_header=variable_header)
         out = publish.to_bytes()
         self.assertEqual(out, b'\x70\x02\x00\x0a')

tests/mqtt/test_pubrec.py

@@ -3,7 +3,7 @@
 # See the file license.txt for copying permission.
 import unittest
 
-from hbmqtt.mqtt.pubrec import PubrecPacket, PubrecVariableHeader
+from hbmqtt.mqtt.pubrec import PubrecPacket, PacketIdVariableHeader
 from hbmqtt.codecs import *
 
 class PubrecPacketTest(unittest.TestCase):
@@ -18,7 +18,7 @@ class PubrecPacketTest(unittest.TestCase):
         self.assertEqual(message.variable_header.packet_id, 10)
 
     def test_to_bytes(self):
-        variable_header = PubrecVariableHeader(10)
+        variable_header = PacketIdVariableHeader(10)
         publish = PubrecPacket(variable_header=variable_header)
         out = publish.to_bytes()
         self.assertEqual(out, b'\x50\x02\x00\x0a')

tests/mqtt/test_pubrel.py

@@ -3,7 +3,7 @@
 # See the file license.txt for copying permission.
 import unittest
 
-from hbmqtt.mqtt.pubrel import PubrelPacket, PubrelVariableHeader
+from hbmqtt.mqtt.pubrel import PubrelPacket, PacketIdVariableHeader
 from hbmqtt.codecs import *
 
 class PubrelPacketTest(unittest.TestCase):
@@ -18,7 +18,7 @@ class PubrelPacketTest(unittest.TestCase):
         self.assertEqual(message.variable_header.packet_id, 10)
 
     def test_to_bytes(self):
-        variable_header = PubrelVariableHeader(10)
+        variable_header = PacketIdVariableHeader(10)
         publish = PubrelPacket(variable_header=variable_header)
         out = publish.to_bytes()
         self.assertEqual(out, b'\x60\x02\x00\x0a')