Omit under attributes, InitVar, ClassVar on dataclasses

Author: ijl

ci/azure-pipelines.yml

@@ -141,7 +141,7 @@ jobs:
   pool:
     vmImage: windows-2019
   variables:
-    interpreter: C:\hostedtoolcache\windows\Python\3.8.2\x64\python.exe
+    interpreter: C:\hostedtoolcache\windows\Python\3.8.3\x64\python.exe
     target: x86_64-pc-windows-msvc
   steps:
   - task: UsePythonVersion@0

pydataclass

@@ -11,6 +11,7 @@ from typing import List
 import orjson
 import rapidjson
 import simplejson
+import ujson
 from tabulate import tabulate
 
 os.sched_setaffinity(os.getpid(), {0, 1})
@@ -63,16 +64,15 @@ table = []
 for lib_name in LIBRARIES:
     if lib_name == "json":
         as_dict = timeit(
-            lambda: json.dumps(objects_as_dict, default=default).encode("utf-8"),
-            number=ITERATIONS,
+            lambda: json.dumps(objects_as_dict).encode("utf-8"), number=ITERATIONS,
         )
         as_dataclass = timeit(
             lambda: json.dumps(objects_as_dataclass, default=default).encode("utf-8"),
             number=ITERATIONS,
         )
     elif lib_name == "simplejson":
         as_dict = timeit(
-            lambda: simplejson.dumps(objects_as_dict, default=default).encode("utf-8"),
+            lambda: simplejson.dumps(objects_as_dict).encode("utf-8"),
             number=ITERATIONS,
         )
         as_dataclass = timeit(
@@ -82,12 +82,13 @@ for lib_name in LIBRARIES:
             number=ITERATIONS,
         )
     elif lib_name == "ujson":
-        as_dict = None
+        as_dict = timeit(
+            lambda: ujson.dumps(objects_as_dict).encode("utf-8"), number=ITERATIONS,
+        )
         as_dataclass = None
     elif lib_name == "rapidjson":
         as_dict = timeit(
-            lambda: rapidjson.dumps(objects_as_dict, default=default).encode("utf-8"),
-            number=ITERATIONS,
+            lambda: rapidjson.dumps(objects_as_dict).encode("utf-8"), number=ITERATIONS,
         )
         as_dataclass = timeit(
             lambda: rapidjson.dumps(objects_as_dataclass, default=default).encode(
@@ -112,7 +113,7 @@ for lib_name in LIBRARIES:
 
     if lib_name == "orjson":
         compared_to_orjson = 1
-    elif as_dict:
+    elif as_dataclass:
         compared_to_orjson = int(as_dataclass / orjson_as_dataclass)
     else:
         compared_to_orjson = None

src/dataclass.rs

@@ -61,12 +61,19 @@ impl<'p> Serialize for DataclassSerializer {
                     std::ptr::null_mut(),
                 )
             };
+            if unsafe { ffi!(PyObject_GetAttr(field, FIELD_TYPE_STR)) != FIELD_TYPE.as_ptr() } {
+                continue;
+            }
             {
                 let data = read_utf8_from_str(attr, &mut str_size);
                 if unlikely!(data.is_null()) {
                     err!(INVALID_STR);
                 }
-                map.serialize_key(str_from_slice!(data, str_size)).unwrap();
+                let key_as_str = str_from_slice!(data, str_size);
+                if key_as_str.as_bytes()[0] == b'_' {
+                    continue;
+                }
+                map.serialize_key(key_as_str).unwrap();
             }
 
             let value = ffi!(PyObject_GetAttr(self.ptr, attr));

src/encode.rs

@@ -239,10 +239,7 @@ impl<'p> Serialize for SerializePyObject {
                 let len = unsafe { PyDict_GET_SIZE(self.ptr) as usize };
                 if unlikely!(len == 0) {
                     serializer.serialize_map(Some(0)).unwrap().end()
-                } else if likely!(
-                    self.opts & SORT_OR_NON_STR_KEYS == 0 || self.opts & DATACLASS_DICT_PATH != 0
-                ) {
-                    let opts = self.opts & !DATACLASS_DICT_PATH;
+                } else if likely!(self.opts & SORT_OR_NON_STR_KEYS == 0) {
                     let mut map = serializer.serialize_map(None).unwrap();
                     let mut pos = 0isize;
                     let mut str_size: pyo3::ffi::Py_ssize_t = 0;
@@ -271,7 +268,7 @@ impl<'p> Serialize for SerializePyObject {
 
                         map.serialize_value(&SerializePyObject::new(
                             value,
-                            opts,
+                            self.opts,
                             self.default_calls,
                             self.recursion + 1,
                             self.default,
@@ -354,15 +351,46 @@ impl<'p> Serialize for SerializePyObject {
                 let dict = ffi!(PyObject_GetAttr(self.ptr, DICT_STR));
                 if !dict.is_null() {
                     ffi!(Py_DECREF(dict));
-                    SerializePyObject::with_obtype(
-                        dict,
-                        ObType::Dict,
-                        self.opts | DATACLASS_DICT_PATH,
-                        self.default_calls,
-                        self.recursion,
-                        self.default,
-                    )
-                    .serialize(serializer)
+                    let len = unsafe { PyDict_GET_SIZE(dict) as usize };
+                    let mut map = serializer.serialize_map(None).unwrap();
+                    let mut pos = 0isize;
+                    let mut str_size: pyo3::ffi::Py_ssize_t = 0;
+                    let mut key: *mut pyo3::ffi::PyObject = std::ptr::null_mut();
+                    let mut value: *mut pyo3::ffi::PyObject = std::ptr::null_mut();
+                    for _ in 0..=len - 1 {
+                        unsafe {
+                            pyo3::ffi::_PyDict_Next(
+                                dict,
+                                &mut pos,
+                                &mut key,
+                                &mut value,
+                                std::ptr::null_mut(),
+                            )
+                        };
+                        if unlikely!(ob_type!(key) != STR_TYPE) {
+                            err!(KEY_MUST_BE_STR)
+                        }
+                        {
+                            let data = read_utf8_from_str(key, &mut str_size);
+                            if unlikely!(data.is_null()) {
+                                err!(INVALID_STR)
+                            }
+                            let key_as_str = str_from_slice!(data, str_size);
+                            if unlikely!(key_as_str.as_bytes()[0] == b'_') {
+                                continue;
+                            }
+                            map.serialize_key(key_as_str).unwrap();
+                        }
+
+                        map.serialize_value(&SerializePyObject::new(
+                            value,
+                            self.opts,
+                            self.default_calls,
+                            self.recursion + 1,
+                            self.default,
+                        ))?;
+                    }
+                    map.end()
                 } else {
                     unsafe { pyo3::ffi::PyErr_Clear() };
                     DataclassSerializer::new(

src/opt.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-pub type Opt = u16;
+pub type Opt = u8;
 
 pub const INDENT_2: Opt = 1;
 pub const NAIVE_UTC: Opt = 1 << 1;
@@ -15,7 +15,6 @@ pub const UTC_Z: Opt = 1 << 7;
 pub const SERIALIZE_DATACLASS: Opt = 0;
 pub const SERIALIZE_UUID: Opt = 0;
 
-pub const DATACLASS_DICT_PATH: Opt = 1 << 8;
 pub const SORT_OR_NON_STR_KEYS: Opt = SORT_KEYS | NON_STR_KEYS;
 
 pub const MAX_OPT: i32 = (INDENT_2

src/typeref.rs

@@ -27,6 +27,7 @@ pub static mut UUID_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 pub static mut ENUM_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 pub static mut ARRAY_TYPE: Lazy<Option<NonNull<PyTypeObject>>> =
     Lazy::new(|| unsafe { look_up_array_type() });
+pub static mut FIELD_TYPE: Lazy<NonNull<PyObject>> = Lazy::new(|| unsafe { look_up_field_type() });
 
 pub static mut BYTES_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 pub static mut BYTEARRAY_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
@@ -38,6 +39,7 @@ pub static mut CONVERT_METHOD_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut DST_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut DICT_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut DATACLASS_FIELDS_STR: *mut PyObject = 0 as *mut PyObject;
+pub static mut FIELD_TYPE_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut ARRAY_STRUCT_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut VALUE_STR: *mut PyObject = 0 as *mut PyObject;
 pub static mut STR_HASH_FUNCTION: Option<hashfunc> = None;
@@ -85,6 +87,7 @@ pub fn init_typerefs() {
         DICT_STR = PyUnicode_InternFromString("__dict__\0".as_ptr() as *const c_char);
         DATACLASS_FIELDS_STR =
             PyUnicode_InternFromString("__dataclass_fields__\0".as_ptr() as *const c_char);
+        FIELD_TYPE_STR = PyUnicode_InternFromString("_field_type\0".as_ptr() as *const c_char);
         ARRAY_STRUCT_STR =
             pyo3::ffi::PyUnicode_InternFromString("__array_struct__\0".as_ptr() as *const c_char);
         VALUE_STR = pyo3::ffi::PyUnicode_InternFromString("value\0".as_ptr() as *const c_char);
@@ -127,6 +130,16 @@ unsafe fn look_up_array_type() -> Option<NonNull<PyTypeObject>> {
     }
 }
 
+unsafe fn look_up_field_type() -> NonNull<PyObject> {
+    let module = PyImport_ImportModule("dataclasses\0".as_ptr() as *const c_char);
+    let module_dict = PyModule_GetDict(module);
+    let ptr = PyMapping_GetItemString(module_dict, "_FIELD\0".as_ptr() as *const c_char)
+        as *mut PyTypeObject;
+    Py_DECREF(module_dict);
+    Py_DECREF(module);
+    NonNull::new_unchecked(ptr as *mut PyObject)
+}
+
 unsafe fn look_up_enum_type() -> *mut PyTypeObject {
     let module = PyImport_ImportModule("enum\0".as_ptr() as *const c_char);
     let module_dict = PyModule_GetDict(module);

test/test_dataclass.py

@@ -2,9 +2,9 @@
 
 import unittest
 import uuid
-from dataclasses import dataclass, field
+from dataclasses import InitVar, dataclass, field
 from enum import Enum
-from typing import Dict, Optional
+from typing import ClassVar, Dict, Optional
 
 import orjson
 
@@ -51,9 +51,12 @@ class Datasubclass(Dataclass1):
 
 @dataclass
 class Slotsdataclass:
-    __slots__ = ("a", "b")
+    __slots__ = ("a", "b", "_c", "d")
     a: str
     b: int
+    _c: str
+    d: InitVar[str]
+    cls_var: ClassVar[str] = "cls"
 
 
 @dataclass
@@ -70,6 +73,18 @@ class UnsortedDataclass:
     d: Optional[Dict]
 
 
+@dataclass
+class InitDataclass:
+    a: InitVar[str]
+    b: InitVar[str]
+    cls_var: ClassVar[str] = "cls"
+    ab: str = ""
+
+    def __post_init__(self, a: str, b: str):
+        self._other = 1
+        self.ab = f"{a} {b}"
+
+
 class DataclassTests(unittest.TestCase):
     def test_dataclass(self):
         """
@@ -146,9 +161,9 @@ def test_dataclass_subclass(self):
 
     def test_dataclass_slots(self):
         """
-        dumps() dataclass with __slots__
+        dumps() dataclass with __slots__ does not include under attributes, InitVar, or ClassVar
         """
-        obj = Slotsdataclass("a", 1)
+        obj = Slotsdataclass("a", 1, "c", "d")
         assert "__dict__" not in dir(obj)
         self.assertEqual(orjson.dumps(obj), b'{"a":"a","b":1}')
 
@@ -191,9 +206,18 @@ def test_dataclass_sort_sub(self):
             b'{"c":1,"b":2,"a":3,"d":{"e":1,"f":2}}',
         )
 
+    def test_dataclass_under(self):
+        """
+        dumps() does not include under attributes, InitVar, or ClassVar
+        """
+        obj = InitDataclass("zxc", "vbn")
+        self.assertEqual(
+            orjson.dumps(obj), b'{"ab":"zxc vbn"}',
+        )
+
     def test_dataclass_option(self):
         """
-        dumps() accepts deprecated OPT_SERIALIZE_DATACALSS
+        dumps() accepts deprecated OPT_SERIALIZE_DATACLASS
         """
         obj = Dataclass1("a", 1, None)
         self.assertEqual(