Catch pygments' exception and raise our own

2 files changed, 15 insertions, 7 deletions

diff --git a/pycco/main.py b/pycco/main.py
index a0e9dfa..f2c20ad 100644
--- a/pycco/main.py
+++ b/pycco/main.py
@@ -347,12 +347,18 @@ def get_language(source, code, language=None):
     if m and m.group(1) in languages:
         return languages[m.group(1)]
     else:
-        lang = lexers.guess_lexer(code).name.lower()
-        for l in languages.values():
-            if l["name"] == lang:
-                return l
-        else:
-            raise ValueError("Can't figure out the language!")
+        try:
+            lang = lexers.guess_lexer(code).name.lower()
+            for l in languages.values():
+                if l["name"] == lang:
+                    return l
+            else:
+                raise ValueError()
+        except ValueError:
+                # If pygments can't find any lexers, it will raise its own
+                # subclass of ValueError. We will catch it and raise ours
+                # for consistency.
+                raise ValueError("Can't figure out the language!")
 
 
 def destination(filepath, preserve_paths=True, outdir=None):
diff --git a/tests/test_pycco.py b/tests/test_pycco.py
index c1a6d7e..1aada4c 100644
--- a/tests/test_pycco.py
+++ b/tests/test_pycco.py
@@ -67,9 +67,11 @@ def test_get_language_bad_source(source):
     code = "#!/usr/bin/python\n"
     code += FOO_FUNCTION
     assert p.get_language(source, code) == PYTHON
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError) as e:
         assert p.get_language(source, "badlang")
 
+    assert e.value.message == "Can't figure out the language!"
+
 
 @given(text() | none())
 def test_get_language_bad_code(code):