Merge pull request #532 from Bedrock02/invalid-po-issue-531

pofile.py: Added new exception called PoFileError and thrown if flagged

2 files changed, 82 insertions, 4 deletions

diff --git a/babel/messages/pofile.py b/babel/messages/pofile.py
index 696ec3e..beb6f35 100644
--- a/babel/messages/pofile.py
+++ b/babel/messages/pofile.py
@@ -19,6 +19,7 @@ from babel.util import wraptext
 from babel._compat import text_type
 
 
+
 def unescape(string):
     r"""Reverse `escape` the given string.
 
@@ -73,6 +74,15 @@ def denormalize(string):
         return unescape(string)
 
 
+class PoFileError(Exception):
+    """Exception thrown by PoParser when an invalid po file is encountered."""
+    def __init__(self, message, catalog, line, lineno):
+        super(PoFileError, self).__init__('{message} on {lineno}'.format(message=message, lineno=lineno))
+        self.catalog = catalog
+        self.line = line
+        self.lineno = lineno
+
+
 class _NormalizedString(object):
 
     def __init__(self, *args):
@@ -104,11 +114,12 @@ class PoFileParser(object):
         'msgid_plural',
     ]
 
-    def __init__(self, catalog, ignore_obsolete=False):
+    def __init__(self, catalog, ignore_obsolete=False, abort_invalid=False):
         self.catalog = catalog
         self.ignore_obsolete = ignore_obsolete
         self.counter = 0
         self.offset = 0
+        self.abort_invalid = abort_invalid
         self._reset_message_state()
 
     def _reset_message_state(self):
@@ -276,11 +287,13 @@ class PoFileParser(object):
             self._add_message()
 
     def _invalid_pofile(self, line, lineno, msg):
+        if self.abort_invalid:
+            raise PoFileError(msg, self.catalog, line, lineno)
         print("WARNING:", msg)
         print("WARNING: Problem on line {0}: {1}".format(lineno + 1, line))
 
 
-def read_po(fileobj, locale=None, domain=None, ignore_obsolete=False, charset=None):
+def read_po(fileobj, locale=None, domain=None, ignore_obsolete=False, charset=None, abort_invalid=False):
     """Read messages from a ``gettext`` PO (portable object) file from the given
     file-like object and return a `Catalog`.
 
@@ -325,9 +338,10 @@ def read_po(fileobj, locale=None, domain=None, ignore_obsolete=False, charset=No
     :param domain: the message domain
     :param ignore_obsolete: whether to ignore obsolete messages in the input
     :param charset: the character set of the catalog.
+    :param abort_invalid: abort read if po file is invalid
     """
     catalog = Catalog(locale=locale, domain=domain, charset=charset)
-    parser = PoFileParser(catalog, ignore_obsolete)
+    parser = PoFileParser(catalog, ignore_obsolete, abort_invalid=abort_invalid)
     parser.parse(fileobj)
     return catalog
 
diff --git a/tests/messages/test_pofile.py b/tests/messages/test_pofile.py
index f6cd66d..002954f 100644
--- a/tests/messages/test_pofile.py
+++ b/tests/messages/test_pofile.py
@@ -13,6 +13,7 @@
 
 from datetime import datetime
 import unittest
+import sys
 
 from babel.core import Locale
 from babel.messages.catalog import Catalog, Message
@@ -20,7 +21,6 @@ from babel.messages import pofile
 from babel.util import FixedOffsetTimezone
 from babel._compat import StringIO, BytesIO
 
-
 class ReadPoTestCase(unittest.TestCase):
 
     def test_preserve_locale(self):
@@ -429,6 +429,70 @@ msgstr[2] "Vohs [text]"
         self.assertEqual("", message.string[1])
         self.assertEqual("Vohs [text]", message.string[2])
 
+    def test_abort_invalid_po_file(self):
+        invalid_po = '''
+            msgctxt ""
+            "{\"checksum\": 2148532640, \"cxt\": \"collector_thankyou\", \"id\": "
+            "270005359}"
+            msgid ""
+            "Thank you very much for your time.\n"
+            "If you have any questions regarding this survey, please contact Fulano "
+            "at nadie@blah.com"
+            msgstr "Merci de prendre le temps de remplir le sondage.
+            Pour toute question, veuillez communiquer avec Fulano  à nadie@blah.com
+            "
+        '''
+        invalid_po_2 = '''
+            msgctxt ""
+            "{\"checksum\": 2148532640, \"cxt\": \"collector_thankyou\", \"id\": "
+            "270005359}"
+            msgid ""
+            "Thank you very much for your time.\n"
+            "If you have any questions regarding this survey, please contact Fulano "
+            "at fulano@blah.com."
+            msgstr "Merci de prendre le temps de remplir le sondage.
+            Pour toute question, veuillez communiquer avec Fulano a fulano@blah.com
+            "
+            '''
+        # Catalog not created, throws Unicode Error
+        buf = StringIO(invalid_po)
+        output = None
+
+        # This should only be thrown under py27
+        if sys.version_info.major == 2:
+            with self.assertRaises(UnicodeEncodeError):
+                output = pofile.read_po(buf, locale='fr', abort_invalid=False)
+            assert not output
+        else:
+            output = pofile.read_po(buf, locale='fr', abort_invalid=False)
+            assert isinstance(output, Catalog)
+
+        # Catalog not created, throws PoFileError
+        buf = StringIO(invalid_po_2)
+        output = None
+        with self.assertRaises(pofile.PoFileError) as e:
+            output = pofile.read_po(buf, locale='fr', abort_invalid=True)
+        assert not output
+
+        # Catalog is created with warning, no abort
+        buf = StringIO(invalid_po_2)
+        output = pofile.read_po(buf, locale='fr', abort_invalid=False)
+        assert isinstance(output, Catalog)
+
+        # Catalog not created, aborted with PoFileError
+        buf = StringIO(invalid_po_2)
+        output = None
+        with self.assertRaises(pofile.PoFileError) as e:
+            output = pofile.read_po(buf, locale='fr', abort_invalid=True)
+        assert not output
+
+    def test_invalid_pofile_with_abort_flag(self):
+        parser = pofile.PoFileParser(None, abort_invalid=True)
+        lineno = 10
+        line = 'Algo esta mal'
+        msg = 'invalid file'
+        with self.assertRaises(pofile.PoFileError) as e:
+            parser._invalid_pofile(line, lineno, msg)
 
 class WritePoTestCase(unittest.TestCase):