mofile: decode context (and optimize with partition)

babel/messages/mofile.py

@@ -62,16 +62,16 @@ def read_mo(fileobj: SupportsRead[bytes]) -> Catalog:
         tlen, toff = unpack(ii, buf[transidx : transidx + 8])
         tend = toff + tlen
         if mend < buflen and tend < buflen:
-            msg = buf[moff:mend]
-            tmsg = buf[toff:tend]
+            raw_msg: bytes = buf[moff:mend]
+            raw_tmsg: bytes = buf[toff:tend]
         else:
             raise OSError(0, 'File is corrupt', filename)
 
         # See if we're looking at GNU .mo conventions for metadata
         if mlen == 0:
             # Catalog description
             lastkey = key = None
-            for item in tmsg.splitlines():
+            for item in raw_tmsg.splitlines():
                 item = item.strip()
                 if not item:
                     continue
@@ -82,19 +82,25 @@ def read_mo(fileobj: SupportsRead[bytes]) -> Catalog:
                 elif lastkey:
                     headers[lastkey] += b'\n' + item
 
-        if b'\x04' in msg:  # context
-            ctxt, msg = msg.split(b'\x04')
+        ctxt: str | None = None
+        msg: list[str] | str
+        tmsg: list[str] | str
+
+        # > Contexts are stored by storing the concatenation of the context,
+        # > a EOT byte, and the original string, instead of the original string.
+        # - https://www.gnu.org/software/gettext/manual/html_node/MO-Files.html
+        ctxt_or_msg, sep, raw_msg = raw_msg.partition(b'\x04')
+        if sep:
+            ctxt = ctxt_or_msg.decode(catalog.charset)
         else:
-            ctxt = None
+            raw_msg = ctxt_or_msg
 
-        if b'\x00' in msg:  # plural forms
-            msg = msg.split(b'\x00')
-            tmsg = tmsg.split(b'\x00')
-            msg = [x.decode(catalog.charset) for x in msg]
-            tmsg = [x.decode(catalog.charset) for x in tmsg]
+        if b'\x00' in raw_msg:  # plural forms
+            msg = [x.decode(catalog.charset) for x in raw_msg.split(b'\x00')]
+            tmsg = [x.decode(catalog.charset) for x in raw_tmsg.split(b'\x00')]
         else:
-            msg = msg.decode(catalog.charset)
-            tmsg = tmsg.decode(catalog.charset)
+            msg = raw_msg.decode(catalog.charset)
+            tmsg = raw_tmsg.decode(catalog.charset)
         catalog[msg] = Message(msg, tmsg, context=ctxt)
 
         # advance to next entry in the seek tables

tests/messages/test_mofile.py

@@ -84,3 +84,23 @@ def test_empty_translation_with_fallback():
     translations.add_fallback(Translations(fp=buf2))
 
     assert translations.ugettext('Fuzz') == 'Flou'
+
+
+def test_read_mo_decodes_message_context() -> None:
+    catalog = Catalog(locale='fi_FI')
+    catalog.add('', '''\
+"Content-Type: text/plain; charset=utf-8\n"
+"Content-Transfer-Encoding: 8bit\n''')
+    catalog.add('x', 'åäöÅÄÖ', context='ほげ')
+    buf = BytesIO()
+    mofile.write_mo(buf, catalog)
+    buf.seek(0)
+    del catalog
+
+    catalog = mofile.read_mo(buf)
+    message = catalog.get('x', context='ほげ')
+    assert message is not None
+    assert message.id == 'x'
+    assert message.context == 'ほげ'
+    assert message.string == 'åäöÅÄÖ'
+    assert not catalog.get('x')  # Not without the context