Improve element not found

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

@@ -10,60 +10,56 @@
 
 @pytest.fixture
 def session():
-
     return get_session()
 
 
 def test_query(session):
-
     si = session.query(Element).filter(Element.symbol == "Si").one()
     assert si.name == "Silicon"
 
 
 def test_element():
-
     si = element("Si")
     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)
     assert e.atomic_number == atomic_number
 
 
 @pytest.mark.parametrize("symbol", SYMBOLS)
 def test_elements_get_by_symbol(symbol):
-
     e = element(symbol)
     assert e.symbol == symbol
 
 
 @pytest.mark.parametrize("name", NAMES)
 def test_elements_get_by_name(name):
-
     e = element(name)
     assert e.name == name
 
 
 @pytest.mark.parametrize("symbol", SYMBOLS)
 def test_elements_str(symbol):
-
     e = element(symbol)
     str(e)
 
 
 @pytest.mark.parametrize("symbol", SYMBOLS)
 def test_elements_repr(symbol):
-
     e = element(symbol)
     repr(e)
 
 
 @pytest.mark.parametrize("symbol", SYMBOLS)
 def test_isotopes_str(symbol):
-
     e = element(symbol)
     for i in e.isotopes:
         str(i)
@@ -72,21 +68,18 @@ def test_isotopes_str(symbol):
 
 @pytest.mark.parametrize("symbol", SYMBOLS)
 def test_electrophilicity(symbol):
-
     e = element(symbol)
     e.electrophilicity()
 
 
 def test__eq__():
-
     elements = get_all_elements()
     for e in elements:
         clone = element(e.symbol)
         assert clone == e
 
 
 def test__ne__():
-
     elements = get_all_elements()
     for e1, e2 in zip(elements[:-1], elements[1:]):
         assert e1 != e2