Added support for infinite and nan floats/doubles

- Added support for the custom double values from
   the protobuf json spec: "Infinity", "-Infinity", and "NaN"
- Added `infinite_floats` test data
- Updated Message.__eq__ to consider nan values
   equal
- Updated `test_message_json` and
   `test_binary_compatibility` to replace NaN float
   values in dictionaries before comparison
   (because two NaN values are not equal)

Author: badge

src/betterproto/__init__.py

@@ -2,6 +2,7 @@
 import enum
 import inspect
 import json
+import math
 import struct
 import sys
 import typing
@@ -117,6 +118,12 @@ def datetime_default_gen() -> datetime:
 DATETIME_ZERO = datetime_default_gen()
 
 
+# Special protobuf json doubles
+INFINITY = "Infinity"
+NEG_INFINITY = "-Infinity"
+NAN = "NaN"
+
+
 class Casing(enum.Enum):
     """Casing constants for serialization."""
 
@@ -373,6 +380,51 @@ def _serialize_single(
     return bytes(output)
 
 
+def _parse_float(value: Any) -> float:
+    """Parse the given value to a float
+
+    Parameters
+    ----------
+    value : Any
+        Value to parse
+
+    Returns
+    -------
+    float
+        Parsed value
+    """
+    if value == INFINITY:
+        return float("inf")
+    if value == NEG_INFINITY:
+        return -float("inf")
+    if value == NAN:
+        return float("nan")
+    return float(value)
+
+
+def _dump_float(value: float) -> Union[float, str]:
+    """Dump the given float to JSON
+
+    Parameters
+    ----------
+    value : float
+        Value to dump
+
+    Returns
+    -------
+    Union[float, str]
+        Dumped valid, either a float or the strings
+        "Infinity" or "-Infinity"
+    """
+    if value == float("inf"):
+        return INFINITY
+    if value == -float("inf"):
+        return NEG_INFINITY
+    if value == float("nan"):
+        return NAN
+    return value
+
+
 def decode_varint(buffer: bytes, pos: int) -> Tuple[int, int]:
     """
     Decode a single varint value from a byte buffer. Returns the value and the
@@ -568,7 +620,18 @@ def __eq__(self, other) -> bool:
                 other_val = other._get_field_default(field_name)
 
             if self_val != other_val:
-                return False
+                # We consider two nan values to be the same for the
+                # purposes of comparing messages (otherwise a message
+                # is not equal to itself)
+                if (
+                    isinstance(self_val, float)
+                    and isinstance(other_val, float)
+                    and math.isnan(self_val)
+                    and math.isnan(other_val)
+                ):
+                    continue
+                else:
+                    return False
 
         return True
 
@@ -1011,6 +1074,11 @@ def to_dict(
                     else:
                         enum_class: Type[Enum] = field_types[field_name]  # noqa
                         output[cased_name] = enum_class(value).name
+                elif meta.proto_type in (TYPE_FLOAT, TYPE_DOUBLE):
+                    if field_is_repeated:
+                        output[cased_name] = [_dump_float(n) for n in value]
+                    else:
+                        output[cased_name] = _dump_float(value)
                 else:
                     output[cased_name] = value
         return output
@@ -1079,6 +1147,11 @@ def from_dict(self: T, value: Dict[str, Any]) -> T:
                             v = [enum_cls.from_string(e) for e in v]
                         elif isinstance(v, str):
                             v = enum_cls.from_string(v)
+                    elif meta.proto_type in (TYPE_FLOAT, TYPE_DOUBLE):
+                        if isinstance(value[key], list):
+                            v = [_parse_float(n) for n in value[key]]
+                        else:
+                            v = _parse_float(value[key])
 
                 if v is not None:
                     setattr(self, field_name, v)

tests/inputs/infinite_floats/infinite_floats.json

@@ -0,0 +1,6 @@
+{
+    "positive": "Infinity",
+    "negative": "-Infinity",
+    "nan": "NaN"
+  }
+  

tests/inputs/infinite_floats/infinite_floats.proto

@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+// Some documentation about the Test message.
+message Test {
+    double positive = 1;
+    double negative = 2;
+    double nan = 3;
+}

tests/test_inputs.py

@@ -1,10 +1,11 @@
 import importlib
 import json
+import math
 import os
 import sys
 from collections import namedtuple
 from types import ModuleType
-from typing import Set
+from typing import Any, Dict, List, Set
 
 import pytest
 
@@ -69,6 +70,56 @@ def module_has_entry_point(module: ModuleType):
     return any(hasattr(module, attr) for attr in ["Test", "TestStub"])
 
 
+def list_replace_nans(l: List) -> List[Any]:
+    """Replace float("nan") in a list with the string "NaN"
+
+    Parameters
+    ----------
+    l : List
+        List to update
+
+    Returns
+    -------
+    List[Any]
+        Updated list
+    """
+    x = []
+    for e in l:
+        if isinstance(e, list):
+            e = list_replace_nans(e)
+        elif isinstance(e, dict):
+            e = dict_replace_nans(e)
+        elif isinstance(e, float) and math.isnan(e):
+            e = betterproto.NAN
+        x.append(e)
+    return x
+
+
+def dict_replace_nans(d: Dict[Any, Any]) -> Dict[Any, Any]:
+    """Replace float("nan") in a dictionary with the string "NaN"
+
+    Parameters
+    ----------
+    l : Dict[Any, Any]
+        Dictionary to update
+
+    Returns
+    -------
+    Dict[Any, Any]
+        Updated dictionary
+    """
+    x = {}
+    for k, v in d.items():
+        if isinstance(v, dict):
+            v = dict_replace_nans(v)
+        elif isinstance(v, list):
+            v = list_replace_nans(v)
+        elif isinstance(v, float) and math.isnan(v):
+            v = betterproto.NAN
+        x[k] = v
+    return x
+
+
 @pytest.fixture
 def test_data(request):
     test_case_name = request.param
@@ -131,7 +182,9 @@ def test_message_json(repeat, test_data: TestData) -> None:
         message.from_json(json_data)
         message_json = message.to_json(0)
 
-        assert json.loads(message_json) == json.loads(json_data)
+        assert dict_replace_nans(json.loads(message_json)) == dict_replace_nans(
+            json.loads(json_data)
+        )
 
 
 @pytest.mark.parametrize("test_data", test_cases.services, indirect=True)
@@ -162,6 +215,6 @@ def test_binary_compatibility(repeat, test_data: TestData) -> None:
         assert bytes(plugin_instance_from_binary) == reference_binary_output
 
         assert plugin_instance_from_json == plugin_instance_from_binary
-        assert (
-            plugin_instance_from_json.to_dict() == plugin_instance_from_binary.to_dict()
-        )
+        assert dict_replace_nans(
+            plugin_instance_from_json.to_dict()
+        ) == dict_replace_nans(plugin_instance_from_binary.to_dict())