Merge pull request #15 from mraspaud/fix-utf-encoding

Make sure messages are consistently encoded as unicode
pytroll · Nov 16, 2018 · 47e25ce · 47e25ce
2 parents 316d87d + 8c6e9cc
commit 47e25ce
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/posttroll/message.py b/posttroll/message.py
@@ -119,7 +119,14 @@ def __init__(self, subject='', atype='', data='', binary=False, rawstr=None):
         if rawstr:
             self.__dict__ = _decode(rawstr)
         else:
-            self.subject = subject
+            try:
+                self.subject = subject.decode('utf-8')
+            except AttributeError:
+                self.subject = subject
+            try:
+                self.type = atype.decode('utf-8')
+            except AttributeError:
+                self.type = atype
             self.type = atype
             self.sender = _getsender()
             self.time = datetime.utcnow()
@@ -168,9 +175,16 @@ def encode(self):
     def __repr__(self):
         return self.encode()
 
-    def __str__(self):
+    def __unicode__(self):
         return self.encode()
 
+    def __str__(self):
+        try:
+            return unicode(self).encode('utf-8')
+        except NameError:
+            return self.encode()
+
+
     def _validate(self):
         """Validate a messages attributes.
         """
@@ -299,7 +313,7 @@ def datetime_encoder(obj):
 def _encode(msg, head=False, binary=False):
     """Convert a Message to a raw string.
     """
-    rawstr = str(_MAGICK) + "{0:s} {1:s} {2:s} {3:s} {4:s}".format(
+    rawstr = str(_MAGICK) + u"{0:s} {1:s} {2:s} {3:s} {4:s}".format(
         msg.subject, msg.type, msg.sender, msg.time.isoformat(), msg.version)
 
     if not head and msg.data:

diff --git a/posttroll/publisher.py b/posttroll/publisher.py
@@ -79,7 +79,7 @@ class Publisher(object):
             while True:
                 counter += 1
                 message = Message("/counter", "info", str(counter))
-                pub.send_string(str(message))
+                pub.send_string(message.encode())
                 time.sleep(3)
         except KeyboardInterrupt:
             print("terminating publisher...")
@@ -247,7 +247,7 @@ class Publish(NoisyPublisher):
                         counter += 1
                         message = Message("/counter", "info", str(counter))
                         print("publishing", message)
-                        pub.send(str(message))
+                        pub.send(message.encode())
                         time.sleep(3)
             except KeyboardInterrupt:
                 print("terminating publisher...")

diff --git a/posttroll/tests/test_pubsub.py b/posttroll/tests/test_pubsub.py
@@ -236,6 +236,24 @@ def test_pub_suber(self):
         self.assertTrue(tested)
         pub.stop()
 
+class TestPub(unittest.TestCase):
+
+    """Testing the publishing capabilities.
+    """
+
+    def setUp(self):
+        test_lock.acquire()
+
+    def tearDown(self):
+        test_lock.release()
+
+    def test_pub_unicode(self):
+        message = Message("/pџтяöll", "info", 'hej')
+        with Publish("a_service", 9000) as pub:
+            try:
+                pub.send(message.encode())
+            except UnicodeDecodeError:
+                self.fail("Sending raised UnicodeDecodeError unexpectedly!")
 
 class TestListenerContainer(unittest.TestCase):
 
@@ -282,5 +300,6 @@ def suite():
     mysuite.addTest(loader.loadTestsFromTestCase(TestNS))
     mysuite.addTest(loader.loadTestsFromTestCase(TestNSWithoutMulticasting))
     mysuite.addTest(loader.loadTestsFromTestCase(TestListenerContainer))
+    mysuite.addTest(loader.loadTestsFromTestCase(TestPub))
 
     return mysuite