Fix Duration export in Result.data and Record.data (#1041)

robsdedude · web-flow · commit dcb337c56fe7 · 2024-04-12T12:21:47.000+02:00
diff --git a/neo4j/data.py b/neo4j/data.py
@@ -206,16 +206,55 @@ def items(self, *keys):
         return list((self.__keys[i], super(Record, self).__getitem__(i)) for i in range(len(self)))
 
     def data(self, *keys):
-        """ Return the keys and values of this record as a dictionary,
-        optionally including only certain values by index or key. Keys
-        provided in the items that are not in the record will be
-        inserted with a value of :const:`None`; indexes provided
-        that are out of bounds will trigger an :exc:`IndexError`.
+        """Return the record as a dictionary.
+
+        Return the keys and values of this record as a dictionary, optionally
+        including only certain values by index or key.
+        Keys provided in the items that are not in the record will be inserted
+        with a value of :data:`None`; indexes provided that are out of bounds
+        will trigger an :exc:`IndexError`.
+
+        This function provides a convenient but opinionated way to transform
+        the record into a mostly JSON serializable format. It is mainly useful
+        for interactive sessions and rapid prototyping.
+
+        The transformation works as follows:
+
+         * Nodes are transformed into dictionaries of their
+           properties.
+
+           * No indication of their original type remains.
+           * Not all information is serialized (e.g., labels and element_id are
+             absent).
+
+         * Relationships are transformed to a tuple of
+           ``(start_node, type, end_node)``, where the nodes are transformed
+           as described above, and type is the relationship type name
+           (:class:`str`).
+
+           * No indication of their original type remains.
+           * No other information (properties, element_id, start_node,
+             end_node, ...) is serialized.
+
+         * Paths are transformed into lists of nodes and relationships. No
+           indication of the original type remains.
+         * :class:`list` and :class:`dict` values are recursively transformed.
+         * Every other type remains unchanged.
+
+           * Spatial types and durations inherit from :class:`tuple`. Hence,
+             they are JSON serializable, but, like graph types, type
+             information will be lost in the process.
+           * The remaining temporal types are not JSON serializable.
+
+        You will have to implement a custom serializer should you need more
+        control over the output format.
+
+        :param keys: Indexes or keys of the items to include. If none are
+            provided, all values will be included.
 
-        :param keys: indexes or keys of the items to include; if none
-                      are provided, all values will be included
         :return: dictionary of values, keyed by field name
-        :raises: :exc:`IndexError` if an out-of-bounds index is specified
+
+        :raises: :exc:`IndexError` if an out-of-bounds index is specified.
         """
         return RecordExporter().transform(dict(self.items(*keys)))
 
@@ -251,7 +290,7 @@ def transform(self, x):
                 path.append(self.transform(relationship.__class__.__name__))
                 path.append(self.transform(x.nodes[i + 1]))
             return path
-        elif isinstance(x, str):
+        elif isinstance(x, (str, Point, Date, Time, DateTime, Duration)):
             return x
         elif isinstance(x, Sequence):
             t = type(x)
diff --git a/neo4j/work/result.py b/neo4j/work/result.py
@@ -366,11 +366,17 @@ def values(self, *keys):
         return [record.values(*keys) for record in self]
 
     def data(self, *keys):
-        """Helper function that return the remainder of the result as a list of dictionaries.
+        """Return the remainder of the result as a list of dictionaries.
 
-        See :class:`neo4j.Record.data`
+        Each dictionary represents a record.
+
+        For details see :meth:`.Record.data`.
+
+        :param keys: Fields to return for each remaining record.
+
+            Optionally filtering to include only certain values by index or
+            key.
 
-        :param keys: fields to return for each remaining record. Optionally filtering to include only certain values by index or key.
         :returns: list of dictionaries
         :rtype: list
         """
diff --git a/tests/unit/test_record.py b/tests/unit/test_record.py
@@ -20,11 +20,24 @@
 
 
 import pytest
+import pytz
 
 from neo4j.data import (
     Graph,
     Node,
+    Path,
     Record,
+    Relationship,
+)
+from neo4j.spatial import (
+    CartesianPoint,
+    WGS84Point,
+)
+from neo4j.time import (
+    Date,
+    DateTime,
+    Duration,
+    Time,
 )
 
 # python -m pytest -s -v tests/unit/test_record.py
@@ -73,17 +86,163 @@ def test_record_repr():
     assert repr(a_record) == "<Record name='Nigel' empire='The British Empire'>"
 
 
-def test_record_data():
-    r = Record(zip(["name", "age", "married"], ["Alice", 33, True]))
-    assert r.data() == {"name": "Alice", "age": 33, "married": True}
-    assert r.data("name") == {"name": "Alice"}
-    assert r.data("age", "name") == {"age": 33, "name": "Alice"}
-    assert r.data("age", "name", "shoe size") == {"age": 33, "name": "Alice", "shoe size": None}
-    assert r.data(0, "name") == {"name": "Alice"}
-    assert r.data(0) == {"name": "Alice"}
-    assert r.data(1, 0) == {"age": 33, "name": "Alice"}
+_RECORD_DATA_ALICE_KEYS = ["name", "age", "married"]
+_RECORD_DATA_ALICE_VALUES = ["Alice", 33, True]
+_RECORD_DATA_GRAPH = Graph()
+_RECORD_DATA_REL_KNOWS = _RECORD_DATA_GRAPH.relationship_type("KNOWS")
+_RECORD_DATA_REL_FOLLOWS = _RECORD_DATA_GRAPH.relationship_type("FOLLOWS")
+
+
+def _record_data_alice_know_bob() -> Relationship:
+    alice = Node(_RECORD_DATA_GRAPH, 1, ["Person"], {"name": "Alice"})
+    bob = Node(_RECORD_DATA_GRAPH, 2, ["Person"], {"name": "Bob"})
+    alice_knows_bob = _RECORD_DATA_REL_KNOWS(_RECORD_DATA_GRAPH, 4,
+                                             {"proper": "tea"})
+    alice_knows_bob._start_node = alice
+    alice_knows_bob._end_node = bob
+    return alice_knows_bob
+
+
+def _record_data_make_path() -> Path:
+    alice_knows_bob = _record_data_alice_know_bob()
+    alice = alice_knows_bob.start_node
+    assert alice is not None
+    bob = alice_knows_bob.end_node
+    carlos = Node(_RECORD_DATA_GRAPH, 3, ["Person"], {"name": "Carlos"})
+    carlos_follows_bob = _RECORD_DATA_REL_FOLLOWS(_RECORD_DATA_GRAPH, 5,
+                                                  {"proper": "tea"})
+    carlos_follows_bob._start_node = carlos
+    carlos_follows_bob._end_node = bob
+    return Path(alice, alice_knows_bob, carlos_follows_bob)
+
+
+@pytest.mark.parametrize(
+    ("keys", "expected"),
+    (
+        (
+            (),
+            {"name": "Alice", "age": 33, "married": True},
+        ),
+        (
+            ("name",),
+            {"name": "Alice"},
+        ),
+        (
+            ("age", "name"),
+            {"age": 33, "name": "Alice"},
+        ),
+        (
+            ("age", "name", "shoe size"),
+            {"age": 33, "name": "Alice", "shoe size": None},
+        ),
+        (
+            ("age", "name", "shoe size"),
+            {"age": 33, "name": "Alice", "shoe size": None},
+        ),
+        (
+            ("age", "name", "shoe size"),
+            {"age": 33, "name": "Alice", "shoe size": None},
+        ),
+        (
+            (0, "name"),
+            {"name": "Alice"},
+        ),
+        (
+            (0,),
+            {"name": "Alice"},
+        ),
+        (
+            (1, 0),
+            {"age": 33, "name": "Alice"},
+        ),
+    ),
+)
+def test_record_data_keys(keys, expected):
+    record = Record(zip(_RECORD_DATA_ALICE_KEYS, _RECORD_DATA_ALICE_VALUES))
+    assert record.data(*keys) == expected
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    (
+        *(
+            (value, value)
+            for value in (
+                None,
+                True,
+                False,
+                0,
+                1,
+                -1,
+                2147483647,
+                -2147483648,
+                3.141592653589,
+                "",
+                "Hello, world!",
+                "👋, 🌍!",
+                [],
+                [1, 2.0, "3", True, None],
+                {"foo": ["bar", 1]},
+                b"",
+                b"foobar",
+                Date(2021, 1, 1),
+                Time(12, 34, 56, 123456789),
+                Time(1, 2, 3, 4, pytz.FixedOffset(60)),
+                DateTime(2021, 1, 1, 12, 34, 56, 123456789),
+                DateTime(2018, 10, 12, 11, 37, 41, 474716862,
+                         pytz.FixedOffset(60)),
+                pytz.timezone("Europe/Stockholm").localize(
+                    DateTime(2018, 10, 12, 11, 37, 41, 474716862)
+                ),
+                Duration(1, 2, 3, 4, 5, 6, 7),
+                CartesianPoint((1, 2.0)),
+                CartesianPoint((1, 2.0, 3)),
+                WGS84Point((1, 2.0)),
+                WGS84Point((1, 2.0, 3)),
+            )
+        ),
+        *(
+            (value, expected)
+            for value, expected in (
+                (
+                    Node(_RECORD_DATA_GRAPH, 1, ["Person"], {"name": "Alice"}),
+                    {"name": "Alice"},
+                ),
+                (
+                    _record_data_alice_know_bob(),
+                    (
+                        {"name": "Alice"},
+                        "KNOWS",
+                        {"name": "Bob"},
+                    )
+                ),
+                (
+                    _record_data_make_path(),
+                    [
+                        {"name": "Alice"},
+                        "KNOWS",
+                        {"name": "Bob"},
+                        "FOLLOWS",
+                        {"name": "Carlos"},
+                    ]
+                ),
+            )
+        ),
+    )
+)
+@pytest.mark.parametrize("wrapper", (None, lambda x: [x], lambda x: {"x": x}))
+def test_record_data_types(value, expected, wrapper):
+    if wrapper is not None:
+        value = wrapper(value)
+        expected = wrapper(expected)
+    record = Record([("key", value)])
+    assert record.data("key") == {"key": expected}
+
+
+def test_record_index_error():
+    record = Record(zip(_RECORD_DATA_ALICE_KEYS, _RECORD_DATA_ALICE_VALUES))
     with pytest.raises(IndexError):
-        _ = r.data(1, 0, 999)
+        record.data(1, 0, 999)
 
 
 def test_record_keys():
@@ -212,13 +371,13 @@ def test_record_get_item():
 
 
 @pytest.mark.parametrize("len_", (0, 1, 2, 42))
-def test_record_len(len_):
+def test_record_len_generic(len_):
     r = Record(("key_%i" % i, "val_%i" % i) for i in range(len_))
     assert len(r) == len_
 
 
 @pytest.mark.parametrize("len_", range(3))
-def test_record_repr(len_):
+def test_record_repr_generic(len_):
     r = Record(("key_%i" % i, "val_%i" % i) for i in range(len_))
     assert repr(r)
 
@@ -275,7 +434,10 @@ def test_record_repr(len_):
         {"x": {"one": 1, "two": 2}}
     ),
     (
-        zip(["a"], [Node("graph", 42, "Person", {"name": "Alice"})]),
+        zip(
+            ["a"],
+            [Node(None, 42, "Person", {"name": "Alice"})]
+        ),
         (),
         {"a": {"name": "Alice"}}
     ),
