Add __id field for record types (#20)

Author: giograno

src/py_avro_schema/_schemas.py

@@ -72,6 +72,7 @@
 NamesType = List[str]
 
 RUNTIME_TYPE_KEY = "_runtime_type"
+REF_ID_KEY = "__id"
 SYMBOL_REGEX = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
 
 
@@ -143,6 +144,9 @@ class Option(enum.Flag):
     #: Adds a _runtime_type field to the record schemas that contains the name of the class
     ADD_RUNTIME_TYPE_FIELD = enum.auto()
 
+    #: Add an __id field to records to track the id of mutable objects
+    ADD_REFERENCE_ID = enum.auto()
+
 
 JSON_OPTIONS = [opt for opt in Option if opt.name and opt.name.startswith("JSON_")]
 
@@ -1008,6 +1012,8 @@ def data_before_deduplication(self, names: NamesType) -> JSONObj:
             doc = _doc_for_class(self.py_type)
             if doc:
                 record_schema["doc"] = doc
+        if Option.ADD_REFERENCE_ID in self.options:
+            record_schema["fields"].append({"name": REF_ID_KEY, "type": ["null", "long"], "default": None})
         return record_schema
 
 

tests/test_dataclass.py

@@ -19,6 +19,7 @@
 import pytest
 
 import py_avro_schema as pas
+from py_avro_schema import Option
 from py_avro_schema._alias import Alias, register_type_aliases
 from py_avro_schema._testing import assert_schema
 
@@ -863,3 +864,26 @@ class PyType:
         ],
     }
     assert_schema(PyType, expected)
+
+
+def test_reference_id():
+    @dataclasses.dataclass
+    class PyType:
+        var: str
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {
+                "name": "var",
+                "type": "string",
+            },
+            {
+                "default": None,
+                "name": "__id",
+                "type": ["null", "long"],
+            },
+        ],
+    }
+    assert_schema(PyType, expected, options=Option.ADD_REFERENCE_ID)

tests/test_plain_class.py

@@ -15,6 +15,7 @@
 import pytest
 
 import py_avro_schema
+from py_avro_schema import Option
 from py_avro_schema._alias import Alias, Opaque, register_type_aliases
 from py_avro_schema._testing import assert_schema
 
@@ -179,3 +180,25 @@ def __init__(self):
     }
 
     assert_schema(PyType, expected)
+
+
+def test_reference_id():
+    class PyType:
+        var: str
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {
+                "name": "var",
+                "type": "string",
+            },
+            {
+                "default": None,
+                "name": "__id",
+                "type": ["null", "long"],
+            },
+        ],
+    }
+    assert_schema(PyType, expected, options=Option.ADD_REFERENCE_ID)

tests/test_typed_dict.py

@@ -131,3 +131,25 @@ class PyType(TypedDict):
     }
 
     assert_schema(PyType, expected, options=pas.Option.MARK_NON_TOTAL_TYPED_DICTS)
+
+
+def test_reference_id():
+    class PyType(TypedDict):
+        var: str
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {
+                "name": "var",
+                "type": "string",
+            },
+            {
+                "default": None,
+                "name": "__id",
+                "type": ["null", "long"],
+            },
+        ],
+    }
+    assert_schema(PyType, expected, options=pas.Option.ADD_REFERENCE_ID)