Integration of orjson for (much) faster JSON encoding and decoding

.pre-commit-config.yaml

@@ -38,6 +38,7 @@ repos:
     rev: 5.10.1
     hooks:
       - id: isort
+        args: ["--profile", "black"]
 
   - repo: https://github.com/psf/black
     rev: 22.1.0

docs_rst/_themes/flask_theme_support.py

@@ -1,7 +1,8 @@
 # flasky extensions.  flasky pygments style based on tango style
 from pygments.style import Style
-from pygments.token import Keyword, Name, Comment, String, Error, \
-     Number, Operator, Generic, Whitespace, Punctuation, Other, Literal
+from pygments.token import (Comment, Error, Generic, Keyword, Literal, Name,
+                            Number, Operator, Other, Punctuation, String,
+                            Whitespace)
 
 
 class FlaskyStyle(Style):
@@ -10,12 +11,12 @@ class FlaskyStyle(Style):
 
     styles = {
         # No corresponding class for the following:
-        #Text:                     "", # class:  ''
+        # Text:                     "", # class:  ''
         Whitespace:                "underline #f8f8f8",      # class: 'w'
-        Error:                     "#a40000 border:#ef2929", # class: 'err'
+        Error:                     "#a40000 border:#ef2929",  # class: 'err'
         Other:                     "#000000",                # class 'x'
 
-        Comment:                   "italic #8f5902", # class: 'c'
+        Comment:                   "italic #8f5902",  # class: 'c'
         Comment.Preproc:           "noitalic",       # class: 'cp'
 
         Keyword:                   "bold #004461",   # class: 'k'
@@ -62,7 +63,7 @@ class FlaskyStyle(Style):
         String:                    "#4e9a06",        # class: 's'
         String.Backtick:           "#4e9a06",        # class: 'sb'
         String.Char:               "#4e9a06",        # class: 'sc'
-        String.Doc:                "italic #8f5902", # class: 'sd' - like a comment
+        String.Doc:                "italic #8f5902",  # class: 'sd' - like a comment
         String.Double:             "#4e9a06",        # class: 's2'
         String.Escape:             "#4e9a06",        # class: 'se'
         String.Heredoc:            "#4e9a06",        # class: 'sh'
@@ -74,7 +75,7 @@ class FlaskyStyle(Style):
 
         Generic:                   "#000000",        # class: 'g'
         Generic.Deleted:           "#a40000",        # class: 'gd'
-        Generic.Emph:              "italic #000000", # class: 'ge'
+        Generic.Emph:              "italic #000000",  # class: 'ge'
         Generic.Error:             "#ef2929",        # class: 'gr'
         Generic.Heading:           "bold #000080",   # class: 'gh'
         Generic.Inserted:          "#00A000",        # class: 'gi'

docs_rst/conf.py

@@ -11,8 +11,8 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
-import sys
 import os
+import sys
 
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
@@ -22,7 +22,7 @@
 sys.path.insert(0, os.path.dirname('../monty'))
 sys.path.insert(0, os.path.dirname('../..'))
 
-from monty import __version__, __author__
+from monty import __author__, __version__
 
 # -- General configuration -----------------------------------------------------
 
@@ -181,7 +181,7 @@
 
 # -- Options for LaTeX output ------------------------------------------------
 
-latex_elements = {
+latex_elements = { # type: ignore
 # The paper size ('letterpaper' or 'a4paper').
 #'papersize': 'letterpaper',
 

monty/design_patterns.py

@@ -116,4 +116,3 @@ def write(*args):  # pylint: disable=E0211
         Does nothing...
         :param args:
         """
-        pass

monty/json.py

@@ -38,6 +38,11 @@
 except ImportError:
     YAML = None  # type: ignore
 
+try:
+    import orjson
+except ImportError:
+    orjson = None  # type: ignore
+
 __version__ = "3.0.0"
 
 
@@ -248,9 +253,7 @@ class MontyEncoder(json.JSONEncoder):
     """
     A Json Encoder which supports the MSONable API, plus adds support for
     numpy arrays, datetime objects, bson ObjectIds (requires bson).
-
     Usage::
-
         # Add it as a *cls* keyword when using json.dump
         json.dumps(object, cls=MontyEncoder)
     """
@@ -262,10 +265,8 @@ def default(self, o) -> dict:  # pylint: disable=E0202
         output. (b) If the @module and @class keys are not in the to_dict,
         add them to the output automatically. If the object has no to_dict
         property, the default Python json encoder default method is called.
-
         Args:
             o: Python object.
-
         Return:
             Python dict representation.
         """
@@ -443,7 +444,10 @@ def decode(self, s):
         :param s: string
         :return: Object.
         """
-        d = json.JSONDecoder.decode(self, s)
+        if orjson is not None:
+            d = orjson.loads(s)  # pylint: disable=E1101
+        else:
+            d = json.loads(s)
         return self.process_decoded(d)
 
 

pylintrc

@@ -368,7 +368,7 @@ ignore-comments=yes
 ignore-docstrings=yes
 
 # Ignore imports when computing similarities.
-ignore-imports=yes
+ignore-imports=no
 
 # Minimum lines number of a similarity.
 min-similarity-lines=4

requirements-ci.txt

@@ -14,4 +14,7 @@ pydantic==1.9.0
 flake8==4.0.1
 pandas==1.4.1
 black==22.1.0
-pylint==2.12.2
+pylint==2.12.2
+orjson==3.6.7
+types-orjson==3.6.2
+types-requests==2.27.11

setup.py

@@ -1,6 +1,6 @@
 import os
-from setuptools import setup, find_packages
-import io
+
+from setuptools import find_packages, setup
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 

tasks.py

@@ -5,17 +5,17 @@
 """
 
 
+import datetime
 import glob
-import requests
 import json
 import os
-import datetime
 import re
 
+import requests
 from invoke import task
-from monty.os import cd
-from monty import __version__ as ver
 
+from monty import __version__ as ver
+from monty.os import cd
 
 __author__ = "Shyue Ping Ong"
 __copyright__ = "Copyright 2012, The Materials Project"

tests/test_bisect.py

@@ -1,6 +1,6 @@
 import unittest
 
-from monty.bisect import index, find_lt, find_le, find_gt, find_ge
+from monty.bisect import find_ge, find_gt, find_le, find_lt, index
 
 
 class FuncTestCase(unittest.TestCase):

tests/test_collections.py

@@ -1,7 +1,7 @@
-import unittest
 import os
+import unittest
 
-from monty.collections import frozendict, Namespace, AttrDict, FrozenAttrDict, tree
+from monty.collections import AttrDict, FrozenAttrDict, Namespace, frozendict, tree
 
 test_dir = os.path.join(os.path.dirname(__file__), "test_files")
 

tests/test_design_patterns.py

@@ -1,7 +1,6 @@
 import unittest
-import pickle
 
-from monty.design_patterns import singleton, cached_class
+from monty.design_patterns import cached_class, singleton
 
 
 class SingletonTest(unittest.TestCase):

tests/test_dev.py

@@ -1,7 +1,8 @@
+import multiprocessing
 import unittest
 import warnings
-import multiprocessing
-from monty.dev import deprecated, requires, get_ncpus, install_excepthook
+
+from monty.dev import deprecated, get_ncpus, install_excepthook, requires
 
 
 class A:

tests/test_fractions.py

@@ -1,6 +1,6 @@
 import unittest
 
-from monty.fractions import gcd, lcm, gcd_float
+from monty.fractions import gcd, gcd_float, lcm
 
 
 class FuncTestCase(unittest.TestCase):

tests/test_functools.py

@@ -1,15 +1,16 @@
-import unittest
-import sys
 import platform
+import sys
 import time
+import unittest
+
 from monty.functools import (
-    lru_cache,
+    TimeoutError,
     lazy_property,
+    lru_cache,
+    prof_main,
     return_if_raise,
     return_none_if_raise,
     timeout,
-    TimeoutError,
-    prof_main,
 )
 
 

tests/test_inspect.py

@@ -1,6 +1,6 @@
 import unittest
 
-from monty.inspect import *
+from monty.inspect import all_subclasses, caller_name, find_top_pyfile
 
 
 class LittleCatA:

tests/test_io.py

@@ -1,17 +1,17 @@
-import unittest
 import os
+import unittest
 
 try:
     from pathlib import Path
 except ImportError:
     Path = None  # type: ignore
 
 from monty.io import (
-    reverse_readline,
-    zopen,
     FileLock,
     FileLockException,
     reverse_readfile,
+    reverse_readline,
+    zopen,
 )
 
 test_dir = os.path.join(os.path.dirname(__file__), "test_files")

tests/test_json.py

@@ -1,17 +1,18 @@
 __version__ = "0.1"
 
+import datetime
+import json
 import os
 import unittest
+from enum import Enum
+
 import numpy as np
-import json
-import datetime
 import pandas as pd
 from bson.objectid import ObjectId
-from enum import Enum
+
+from monty.json import MontyDecoder, MontyEncoder, MSONable, _load_redirect, jsanitize
 
 from . import __version__ as tests_version
-from monty.json import MSONable, MontyEncoder, MontyDecoder, jsanitize
-from monty.json import _load_redirect
 
 test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files")
 
@@ -101,6 +102,11 @@ def __init__(self, s):
         self.s = s
 
 
+class ClassContainingNumpyArray(MSONable):
+    def __init__(self, np_a):
+        self.np_a = np_a
+
+
 class MSONableTest(unittest.TestCase):
     def setUp(self):
         self.good_cls = GoodMSONClass
@@ -146,7 +152,7 @@ def test_to_from_dict(self):
         self.assertRaises(NotImplementedError, obj.as_dict)
         obj = self.auto_mson(2, 3)
         d = obj.as_dict()
-        objd = self.auto_mson.from_dict(d)
+        self.auto_mson.from_dict(d)
 
     def test_unsafe_hash(self):
         GMC = GoodMSONClass
@@ -288,7 +294,7 @@ def test_datetime(self):
         self.assertEqual(type(d["dt"]), datetime.datetime)
 
     def test_uuid(self):
-        from uuid import uuid4, UUID
+        from uuid import UUID, uuid4
 
         uuid = uuid4()
         jsonstr = json.dumps(uuid, cls=MontyEncoder)
@@ -343,6 +349,22 @@ def test_numpy(self):
         d = jsanitize(x, strict=True)
         assert type(d["energies"][0]) == float
 
+        # Test data nested in a class
+        x = np.array([[1 + 1j, 2 + 1j], [3 + 1j, 4 + 1j]], dtype="complex64")
+        cls = ClassContainingNumpyArray(np_a={"a": [{"b": x}]})
+
+        d = json.loads(json.dumps(cls, cls=MontyEncoder))
+
+        self.assertEqual(d["np_a"]["a"][0]["b"]["@module"], "numpy")
+        self.assertEqual(d["np_a"]["a"][0]["b"]["@class"], "array")
+        self.assertEqual(d["np_a"]["a"][0]["b"]["data"], [[[1.0, 2.0], [3.0, 4.0]], [[1.0, 1.0], [1.0, 1.0]]])
+        self.assertEqual(d["np_a"]["a"][0]["b"]["dtype"], "complex64")
+
+        obj = ClassContainingNumpyArray.from_dict(d)
+        self.assertIsInstance(obj, ClassContainingNumpyArray)
+        self.assertIsInstance(obj.np_a["a"][0]["b"], np.ndarray)
+        self.assertEqual(obj.np_a["a"][0]["b"][0][1], 2 + 1j)
+
     def test_pandas(self):
 
         cls = ClassContainingDataFrame(df=pd.DataFrame([{"a": 1, "b": 1}, {"a": 1, "b": 2}]))
@@ -369,6 +391,18 @@ def test_pandas(self):
         self.assertIsInstance(obj.s, pd.Series)
         self.assertEqual(list(obj.s.a), [1, 2, 3])
 
+        cls = ClassContainingSeries(s={"df": [pd.Series({"a": [1, 2, 3], "b": [4, 5, 6]})]})
+
+        d = json.loads(MontyEncoder().encode(cls))
+
+        self.assertEqual(d["s"]["df"][0]["@module"], "pandas")
+        self.assertEqual(d["s"]["df"][0]["@class"], "Series")
+
+        obj = ClassContainingSeries.from_dict(d)
+        self.assertIsInstance(obj, ClassContainingSeries)
+        self.assertIsInstance(obj.s["df"][0], pd.Series)
+        self.assertEqual(list(obj.s["df"][0].a), [1, 2, 3])
+
     def test_callable(self):
         instance = MethodSerializationClass(a=1)
         for function in [

tests/test_logging.py

@@ -1,7 +1,7 @@
+import logging
 import unittest
 from io import StringIO
 
-import logging
 from monty.logging import logged
 
 

tests/test_multiprocessing.py

@@ -1,7 +1,7 @@
 import unittest
+from math import sqrt
 
 from monty.multiprocessing import imap_tqdm
-from math import sqrt
 
 
 class FuncCase(unittest.TestCase):

tests/test_operator.py

@@ -1,4 +1,5 @@
 import unittest
+
 from monty.operator import operator_from_str