Fixup YapfError messages

4 files changed, 16 insertions, 8 deletions

diff --git a/yapf/__init__.py b/yapf/__init__.py
index d9b4dd6..d09d4c7 100644
--- a/yapf/__init__.py
+++ b/yapf/__init__.py
@@ -103,6 +103,8 @@ def main(argv):
           style_config=style_config,
           lines=lines,
           verify=args.verify)
+    except errors.YapfError:
+      raise
     except Exception as e:
       raise errors.YapfError(errors.FormatErrorMsg(e))
 
diff --git a/yapf/yapflib/errors.py b/yapf/yapflib/errors.py
index 5726ff1..99e88d9 100644
--- a/yapf/yapflib/errors.py
+++ b/yapf/yapflib/errors.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 """YAPF error objects."""
 
+from lib2to3.pgen2 import tokenize
+
 
 def FormatErrorMsg(e):
   """Convert an exception into a standard format.
@@ -27,6 +29,11 @@ def FormatErrorMsg(e):
   Returns:
     A properly formatted error message string.
   """
+  if isinstance(e, SyntaxError):
+    return '{}:{}:{}: {}'.format(e.filename, e.lineno, e.offset, e.msg)
+  if isinstance(e, tokenize.TokenError):
+    return '{}:{}:{}: {}'.format(e.filename, e.args[1][0], e.args[1][1],
+                                 e.args[0])
   return '{}:{}:{}: {}'.format(e.args[1][0], e.args[1][1], e.args[1][2], e.msg)
 
 
diff --git a/yapf/yapflib/yapf_api.py b/yapf/yapflib/yapf_api.py
index ac0466e..06c9a73 100644
--- a/yapf/yapflib/yapf_api.py
+++ b/yapf/yapflib/yapf_api.py
@@ -181,13 +181,9 @@ def FormatCode(unformatted_source,
   """
   try:
     tree = pytree_utils.ParseCodeToTree(unformatted_source)
-  except tokenize.TokenError as e:
-    e.msg = e.args[0]
-    e.args = (e.msg, (filename, e.args[1][0], e.args[1][1]))
-    raise
   except Exception as e:
-    e.args = (e.args[0], (filename, e.args[1][1], e.args[1][2], e.args[1][3]))
-    raise
+    e.filename = filename
+    raise errors.YapfError(errors.FormatErrorMsg(e))
 
   reformatted_source = FormatTree(
       tree, style_config=style_config, lines=lines, verify=verify)
diff --git a/yapftests/yapf_test.py b/yapftests/yapf_test.py
index e116043..2330f4e 100644
--- a/yapftests/yapf_test.py
+++ b/yapftests/yapf_test.py
@@ -26,6 +26,7 @@ import unittest
 
 from lib2to3.pgen2 import tokenize
 
+from yapf.yapflib import errors
 from yapf.yapflib import py3compat
 from yapf.yapflib import style
 from yapf.yapflib import yapf_api
@@ -62,6 +63,7 @@ class FormatCodeTest(yapf_test_helper.YAPFTest):
         """)
     self._Check(unformatted_code, expected_formatted_code)
 
+  @unittest.skipUnless(py3compat.PY36, 'Requires Python 3.6')
   def testPrintAfterPeriod(self):
     unformatted_code = textwrap.dedent("""a.print\n""")
     expected_formatted_code = textwrap.dedent("""a.print\n""")
@@ -428,6 +430,7 @@ class CommandLineTest(unittest.TestCase):
     self.assertEqual(stderrdata, b'')
     self.assertMultiLineEqual(reformatted_code.decode('utf-8'), expected)
 
+  @unittest.skipUnless(py3compat.PY36, 'Requires Python 3.6')
   def testUnicodeEncodingPipedToFile(self):
     unformatted_code = textwrap.dedent(u"""\
         def foo():
@@ -1560,11 +1563,11 @@ class BadInputTest(unittest.TestCase):
 
   def testBadSyntax(self):
     code = '  a = 1\n'
-    self.assertRaises(SyntaxError, yapf_api.FormatCode, code)
+    self.assertRaises(errors.YapfError, yapf_api.FormatCode, code)
 
   def testBadCode(self):
     code = 'x = """hello\n'
-    self.assertRaises(tokenize.TokenError, yapf_api.FormatCode, code)
+    self.assertRaises(errors.YapfError, yapf_api.FormatCode, code)
 
 
 class DiffIndentTest(unittest.TestCase):