Improve element not found (#142)

* Improve element not found

* Fix formatting

mendeleev/mendeleev.py

@@ -2,6 +2,8 @@
 
 from typing import List, Union
 
+import sqlalchemy
+
 from .db import get_session
 from .models import Element, Isotope
 
@@ -74,15 +76,17 @@ def _get_element(ids) -> Union[Element, List[Element]]:
 
     session = get_session()
 
-    if isinstance(ids, str):
-        if len(ids) <= 3 and ids.lower() != "tin":
-            return session.query(Element).filter(Element.symbol == str(ids)).one()
-        else:
-            return session.query(Element).filter(Element.name == str(ids)).one()
-    elif isinstance(ids, int):
-        return session.query(Element).filter(Element.atomic_number == ids).one()
-    else:
+    try:
+        if isinstance(ids, str):
+            if len(ids) <= 3 and ids.lower() != "tin":
+                return session.query(Element).filter(Element.symbol == str(ids)).one()
+            else:
+                return session.query(Element).filter(Element.name == str(ids)).one()
+        elif isinstance(ids, int):
+            return session.query(Element).filter(Element.atomic_number == ids).one()
         raise ValueError("Expecting a <str> or <int>, got: {0:s}".format(type(ids)))
+    except sqlalchemy.exc.NoResultFound:
+        raise ValueError(f"Element not found: {ids}")
 
 
 def get_all_elements() -> List[Element]:

tests/test_element.py

@@ -23,6 +23,11 @@ def test_element():
     assert si.name == "Silicon"
 
 
+def test_incorrect_element():
+    with pytest.raises(ValueError, match="Element not found: si"):
+        element("si")
+
+
 @pytest.mark.parametrize("atomic_number", list(range(1, 119)))
 def test_elements_get_by_atomic_number(atomic_number):
     e = element(atomic_number)