diff --git a/src/monty/itertools.py b/src/monty/itertools.py index 5a112a03..c4ed76d5 100644 --- a/src/monty/itertools.py +++ b/src/monty/itertools.py @@ -7,10 +7,7 @@ import itertools from typing import TYPE_CHECKING -try: - import numpy as np -except ImportError: - np = None +import numpy as np if TYPE_CHECKING: from typing import Iterable diff --git a/src/monty/json.py b/src/monty/json.py index 282cfad3..f03a7168 100644 --- a/src/monty/json.py +++ b/src/monty/json.py @@ -21,13 +21,9 @@ from typing import Any from uuid import UUID, uuid4 +import numpy as np from ruamel.yaml import YAML -try: - import numpy as np -except ImportError: - np = None - try: import pydantic except ImportError: @@ -595,23 +591,22 @@ def default(self, o) -> dict: d["data"] = o.numpy().tolist() return d - if np is not None: - if isinstance(o, np.ndarray): - if str(o.dtype).startswith("complex"): - return { - "@module": "numpy", - "@class": "array", - "dtype": str(o.dtype), - "data": [o.real.tolist(), o.imag.tolist()], - } + if isinstance(o, np.ndarray): + if str(o.dtype).startswith("complex"): return { "@module": "numpy", "@class": "array", "dtype": str(o.dtype), - "data": o.tolist(), + "data": [o.real.tolist(), o.imag.tolist()], } - if isinstance(o, np.generic): - return o.item() + return { + "@module": "numpy", + "@class": "array", + "dtype": str(o.dtype), + "data": o.tolist(), + } + if isinstance(o, np.generic): + return o.item() if _check_type(o, "pandas.core.frame.DataFrame"): return { @@ -809,7 +804,7 @@ def process_decoded(self, d): ).type(d["dtype"]) return torch.tensor(d["data"]).type(d["dtype"]) # pylint: disable=E1101 - elif np is not None and modname == "numpy" and classname == "array": + elif modname == "numpy" and classname == "array": if d["dtype"].startswith("complex"): return np.array( [ @@ -932,7 +927,8 @@ def jsanitize( ) for i in obj ] - if np is not None and isinstance(obj, np.ndarray): + + if isinstance(obj, np.ndarray): try: return [ jsanitize( @@ -946,8 +942,10 @@ def jsanitize( ] except TypeError: return obj.tolist() - if np is not None and isinstance(obj, np.generic): + + if isinstance(obj, np.generic): return obj.item() + if _check_type( obj, ( diff --git a/tests/test_json.py b/tests/test_json.py index 9c9e572c..097f1f53 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -8,10 +8,7 @@ from enum import Enum from typing import Union -try: - import numpy as np -except ImportError: - np = None +import numpy as np try: import pandas as pd @@ -564,7 +561,6 @@ def test_nan(self): d = json.loads(djson) assert isinstance(d[0], float) - @pytest.mark.skipif(np is None, reason="numpy not present") def test_numpy(self): x = np.array([1, 2, 3], dtype="int64") with pytest.raises(TypeError): @@ -872,9 +868,7 @@ def test_jsanitize_pandas(self): clean = jsanitize(s) assert clean == s.to_dict() - @pytest.mark.skipif( - np is None or ObjectId is None, reason="numpy and bson not present" - ) + @pytest.mark.skipif(ObjectId is None, reason="bson not present") def test_jsanitize_numpy_bson(self): d = { "a": ["b", np.array([1, 2, 3])],