Type annotations for canonicaljson

Author: David Robertson

MANIFEST.in

@@ -3,5 +3,6 @@ include *.py
 include *.md
 include LICENSE
 include tox.ini
+include pyproject.toml
 prune .travis
 prune debian

canonicaljson.py

@@ -16,9 +16,14 @@
 # limitations under the License.
 
 import platform
-from typing import Optional, Type
+from typing import Any, Generator, Optional, Type
 
-frozendict_type: Optional[Type]
+try:
+    from typing import Protocol
+except ImportError:  # pragma: no cover
+    from typing_extensions import Protocol  # type: ignore[misc]
+
+frozendict_type: Optional[Type[Any]]
 try:
     from frozendict import frozendict as frozendict_type
 except ImportError:
@@ -27,22 +32,37 @@
 __version__ = "1.6.0"
 
 
-def _default(obj):  # pragma: no cover
+def _default(obj: object) -> object:  # pragma: no cover
     if type(obj) is frozendict_type:
         # If frozendict is available and used, cast `obj` into a dict
-        return dict(obj)
+        return dict(obj)  # type: ignore[call-overload]
     raise TypeError(
         "Object of type %s is not JSON serializable" % obj.__class__.__name__
     )
 
 
+class Encoder(Protocol):  # pragma: no cover
+    def encode(self, data: object) -> str:
+        pass
+
+    def iterencode(self, data: object) -> Generator[str, None, None]:
+        pass
+
+    def __call__(self, *args: Any, **kwargs: Any) -> "Encoder":
+        pass
+
+
+class JsonLibrary(Protocol):
+    JSONEncoder: Encoder
+
+
 # Declare these in the module scope, but they get configured in
 # set_json_library.
-_canonical_encoder = None
-_pretty_encoder = None
+_canonical_encoder: Encoder = None  # type: ignore[assignment]
+_pretty_encoder: Encoder = None   # type: ignore[assignment]
 
 
-def set_json_library(json_lib):
+def set_json_library(json_lib: JsonLibrary) -> None:
     """
     Set the underlying JSON library that canonicaljson uses to json_lib.
 
@@ -69,7 +89,7 @@ def set_json_library(json_lib):
     )
 
 
-def encode_canonical_json(json_object):
+def encode_canonical_json(json_object: object) -> bytes:
     """Encodes the shortest UTF-8 JSON encoding with dictionary keys
     lexicographically sorted by unicode code point.
 
@@ -82,7 +102,7 @@ def encode_canonical_json(json_object):
     return s.encode("utf-8")
 
 
-def iterencode_canonical_json(json_object):
+def iterencode_canonical_json(json_object: object) -> Generator[bytes, None, None]:
     """Encodes the shortest UTF-8 JSON encoding with dictionary keys
     lexicographically sorted by unicode code point.
 
@@ -95,7 +115,7 @@ def iterencode_canonical_json(json_object):
         yield chunk.encode("utf-8")
 
 
-def encode_pretty_printed_json(json_object):
+def encode_pretty_printed_json(json_object: object) -> bytes:
     """
     Encodes the JSON object dict as human readable UTF-8 bytes.
 
@@ -108,7 +128,7 @@ def encode_pretty_printed_json(json_object):
     return _pretty_encoder.encode(json_object).encode("utf-8")
 
 
-def iterencode_pretty_printed_json(json_object):
+def iterencode_pretty_printed_json(json_object: object) -> Generator[bytes, None, None]:
     """Encodes the JSON object dict as human readable UTF-8 bytes.
 
     Args:
@@ -132,7 +152,7 @@ def iterencode_pretty_printed_json(json_object):
     #
     # Note that it seems performance is on par or better using json from the
     # standard library as of Python 3.7.
-    import simplejson as json
+    import simplejson as json  # type: ignore[no-redef]
 
 # Set the JSON library to the backwards compatible version.
 set_json_library(json)

pyproject.toml

@@ -0,0 +1,18 @@
+[tool.mypy]
+show_error_codes = true
+strict = true
+
+files = ["."]
+exclude = "setup.py"
+#mypy_path = "stubs"
+
+#[[tool.mypy.overrides]]
+#module = [
+#    "idna",
+#    "netaddr",
+#    "prometheus_client",
+#    "signedjson.*",
+#    "sortedcontainers",
+#]
+#ignore_missing_imports = true
+

setup.py

@@ -49,6 +49,8 @@ def exec_file(path_segments, name):
         # simplerjson versions before 3.14.0 had a bug with some characters
         # (e.g. \u2028) if ensure_ascii was set to false.
         "simplejson>=3.14.0",
+        # typing.Protocol was only added to the stdlib in Python 3.8
+        "typing_extensions>=4.0.0; python_version < '3.8'",
     ],
     extras_require={
         # frozendict support can be enabled using the `canonicaljson[frozendict]` syntax

test_canonicaljson.py

@@ -31,7 +31,7 @@
 
 
 class TestCanonicalJson(unittest.TestCase):
-    def test_encode_canonical(self):
+    def test_encode_canonical(self) -> None:
         self.assertEqual(encode_canonical_json({}), b"{}")
 
         # ctrl-chars should be encoded.
@@ -68,7 +68,7 @@ def test_encode_canonical(self):
         # Iteratively encoding should work.
         self.assertEqual(list(iterencode_canonical_json({})), [b"{}"])
 
-    def test_ascii(self):
+    def test_ascii(self) -> None:
         """
         Ensure the proper ASCII characters are escaped.
 
@@ -95,10 +95,10 @@ def test_ascii(self):
         # And other characters are passed unescaped.
         unescaped = [0x20, 0x21] + list(range(0x23, 0x5C)) + list(range(0x5D, 0x7E))
         for c in unescaped:
-            c = chr(c)
-            self.assertEqual(encode_canonical_json(c), b'"' + c.encode("ascii") + b'"')
+            s = chr(c)
+            self.assertEqual(encode_canonical_json(s), b'"' + s.encode("ascii") + b'"')
 
-    def test_encode_pretty_printed(self):
+    def test_encode_pretty_printed(self) -> None:
         self.assertEqual(encode_pretty_printed_json({}), b"{}")
         self.assertEqual(list(iterencode_pretty_printed_json({})), [b"{}"])
 
@@ -112,7 +112,9 @@ def test_encode_pretty_printed(self):
         frozendict_type is None,
         "If `frozendict` is not available, skip test",
     )
-    def test_frozen_dict(self):
+    def test_frozen_dict(self) -> None:
+        # For mypy's benefit:
+        assert frozendict_type is not None
         self.assertEqual(
             encode_canonical_json(frozendict_type({"a": 1})),
             b'{"a":1}',
@@ -122,7 +124,7 @@ def test_frozen_dict(self):
             b'{\n    "a": 1\n}',
         )
 
-    def test_unknown_type(self):
+    def test_unknown_type(self) -> None:
         class Unknown(object):
             pass
 
@@ -133,7 +135,7 @@ class Unknown(object):
         with self.assertRaises(Exception):
             encode_pretty_printed_json(unknown_object)
 
-    def test_invalid_float_values(self):
+    def test_invalid_float_values(self) -> None:
         """Infinity/-Infinity/NaN are not allowed in canonicaljson."""
 
         with self.assertRaises(ValueError):
@@ -154,7 +156,7 @@ def test_invalid_float_values(self):
         with self.assertRaises(ValueError):
             encode_pretty_printed_json(nan)
 
-    def test_set_json(self):
+    def test_set_json(self) -> None:
         """Ensure that changing the underlying JSON implementation works."""
         mock_json = mock.Mock(spec=["JSONEncoder"])
         mock_json.JSONEncoder.return_value.encode.return_value = "sentinel"
@@ -163,6 +165,6 @@ def test_set_json(self):
             self.assertEqual(encode_canonical_json({}), b"sentinel")
         finally:
             # Reset the JSON library to whatever was originally set.
-            from canonicaljson import json
+            from canonicaljson import json  # type: ignore[attr-defined]
 
             set_json_library(json)

tox.ini

@@ -26,3 +26,11 @@ basepython = python3.7
 deps =
     black==21.9b0
 commands = python -m black --check --diff .
+
+[testenv:mypy]
+deps =
+    mypy==0.942
+    types-frozendict==2.0.8
+    types-simplejson==3.17.5
+    types-setuptools==57.4.14
+commands = mypy