[CLIENT-2417] Replace partitions status dictionary with aerospike.PartitionsStatus class

aerospike-stubs/aerospike.pyi

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Union, final, Literal, Optional, Final
+from typing import Any, Callable, Union, final, Literal, Optional
 
 from aerospike_helpers.batch.records import BatchRecords
 from aerospike_helpers.metrics import MetricsPolicy
@@ -463,3 +463,9 @@ def set_log_handler(callback: Callable = ...) -> None: ...
 def set_log_level(log_level: int) -> None: ...
 def set_serializer(callback: Callable) -> None: ...
 def unset_serializers() -> None: ...
+
+class PartitionStatus:
+    def __getitem__(self, attr): ...
+
+class PartitionsStatus:
+    def __getitem__(self, attr): ...

src/include/conversions.h

@@ -209,16 +209,11 @@ as_status convert_exp_list(AerospikeClient *self, PyObject *py_exp_list,
 as_status convert_partition_filter(AerospikeClient *self,
                                    PyObject *py_partition_filter,
                                    as_partition_filter *partition_filter,
-                                   as_partitions_status **ps, as_error *err);
+                                   as_error *err);
 
 as_status get_int_from_py_int(as_error *err, PyObject *py_long,
                               int *int_pointer, const char *py_object_name);
 
-as_status
-as_partitions_status_to_pyobject(as_error *err,
-                                 const as_partitions_status *parts_status,
-                                 PyObject **py_dict);
-
 as_status as_partition_status_to_pyobject(
     as_error *err, const as_partition_status *part_status, PyObject **py_tuple);
 

src/include/partitions_status.h

@@ -0,0 +1,11 @@
+#include <Python.h>
+
+#include <aerospike/as_partition_filter.h>
+
+PyTypeObject *AerospikePartitionsStatusObject_Ready();
+PyTypeObject *AerospikePartitionStatusObject_Ready();
+
+PyObject *create_py_partitions_status_object(as_error *err,
+                                             as_partitions_status *parts_all);
+
+extern PyTypeObject AerospikePartitionsStatusObject_Type;

src/include/types.h

@@ -103,6 +103,14 @@ typedef struct {
     PyDictObject dict;
 } AerospikeKeyOrderedDict;
 
+typedef struct {
+    PyObject_HEAD as_partitions_status *parts_all;
+} AerospikePartitionsStatusObject;
+
+typedef struct {
+    PyObject_HEAD as_partition_status *part_status;
+} AerospikePartitionStatusObject;
+
 typedef struct {
     PyObject_HEAD
         /* Type-specific fields go here. */

src/main/aerospike.c

@@ -32,6 +32,7 @@
 #include "module_functions.h"
 #include "nullobject.h"
 #include "cdt_types.h"
+#include "partitions_status.h"
 #include "transaction.h"
 #include "config_provider.h"
 
@@ -556,6 +557,8 @@ static struct type_name_to_creation_method py_module_types[] = {
     {"null", AerospikeNullObject_Ready},
     {"CDTWildcard", AerospikeWildcardObject_Ready},
     {"CDTInfinite", AerospikeInfiniteObject_Ready},
+    {"PartitionsStatus", AerospikePartitionsStatusObject_Ready},
+    {"PartitionStatus", AerospikePartitionStatusObject_Ready},
     {"Transaction", AerospikeTransaction_Ready},
     {"ConfigProvider", AerospikeConfigProvider_Ready},
 };

src/main/conversions.c

@@ -356,65 +356,6 @@ as_status as_partition_status_to_pyobject(
     return err->code;
 }
 
-// creates a python dict of tuples from an as_partitions_status
-// EX: {id:(id, init, done, digest, bval) for id in range (1000, 1004,1)}
-// returns and empty dict if parts_status == NULL
-as_status as_partitions_status_to_pyobject(
-    as_error *err, const as_partitions_status *parts_status, PyObject **py_dict)
-{
-    as_error_reset(err);
-
-    PyObject *new_dict = PyDict_New();
-    if (new_dict == NULL) {
-        as_error_update(err, AEROSPIKE_ERR_CLIENT, "failed to create new_dict");
-        goto END;
-    }
-
-    if (parts_status == NULL) {
-        // If parts_status is NULL return an empty dict because
-        // the query/scan is not tracking its partitions.
-        *py_dict = new_dict;
-        goto END;
-    }
-
-    PyObject *py_done = PyBool_FromLong(parts_status->done);
-    PyDict_SetItemString(new_dict, PARTITIONS_STATUS_KEY_DONE, py_done);
-    Py_DECREF(py_done);
-
-    PyObject *py_retry = PyBool_FromLong(parts_status->retry);
-    PyDict_SetItemString(new_dict, PARTITIONS_STATUS_KEY_RETRY, py_retry);
-    Py_DECREF(py_retry);
-
-    for (int i = 0; i < parts_status->part_count; ++i) {
-
-        const as_partition_status *part = &parts_status->parts[i];
-
-        PyObject *new_py_tuple = NULL;
-        if (as_partition_status_to_pyobject(err, part, &new_py_tuple) !=
-            AEROSPIKE_OK) {
-            Py_DECREF(new_dict);
-            goto END;
-        }
-
-        PyObject *py_id = PyLong_FromUnsignedLong((unsigned long)part->part_id);
-
-        if (PyDict_SetItem(new_dict, py_id, new_py_tuple) != 0) {
-            as_error_update(err, AEROSPIKE_ERR_CLIENT,
-                            "failed set item in new_dict");
-            Py_DECREF(new_dict);
-            Py_DECREF(new_py_tuple);
-            Py_XDECREF(py_id);
-            goto END;
-        }
-        Py_DECREF(py_id);
-    }
-
-    *py_dict = new_dict;
-
-END:
-    return err->code;
-}
-
 as_status as_user_info_to_pyobject(as_error *err, as_user *user,
                                    PyObject **py_as_user)
 {