pygeos · caspervdw · Dec 8, 2020 · Nov 12, 2020 · Nov 20, 2020 · Nov 20, 2020
diff --git a/setup.py b/setup.py
@@ -154,6 +154,7 @@ def finalize_options(self):
                 "src/pygeom.c",
                 "src/strtree.c",
                 "src/ufuncs.c",
+                "src/vector.c"
             ],
             **ext_options,
         )

diff --git a/src/strtree.c b/src/strtree.c
@@ -15,6 +15,7 @@
 #include "kvec.h"
 #include "pygeom.h"
 #include "strtree.h"
+#include "vector.h"
 
 /* GEOS function that takes a prepared geometry and a regular geometry
  * and returns bool value */
@@ -61,34 +62,6 @@ FuncGEOS_YpY_b* get_predicate_func(int predicate_id) {
   }
 }
 
-/* Copy values from arr to a new numpy integer array.
- *
- * Parameters
- * ----------
- * arr: dynamic vector array to convert to ndarray
- */
-
-static PyArrayObject* copy_kvec_to_npy(npy_intp_vec* arr) {
-  npy_intp i;
-  npy_intp size = kv_size(*arr);
-
-  npy_intp dims[1] = {size};
-  // the following raises a compiler warning based on how the macro is defined
-  // in numpy.  There doesn't appear to be anything we can do to avoid it.
-  PyArrayObject* result = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_INTP);
-  if (result == NULL) {
-    PyErr_SetString(PyExc_RuntimeError, "could not allocate numpy array");
-    return NULL;
-  }
-
-  for (i = 0; i < size; i++) {
-    // assign value into numpy array
-    *(npy_intp*)PyArray_GETPTR1(result, i) = kv_A(*arr, i);
-  }
-
-  return (PyArrayObject*)result;
-}
-
 static void STRtree_dealloc(STRtreeObject* self) {
   size_t i, size;
 
@@ -114,7 +87,7 @@ static PyObject* STRtree_new(PyTypeObject* type, PyObject* args, PyObject* kwds)
   void *tree, *ptr;
   npy_intp n, i, count = 0;
   GEOSGeometry* geom;
-  pg_geom_obj_vec _geoms;
+  geom_obj_vec_t _geoms;
   GeometryObject* obj;
 
   if (!PyArg_ParseTuple(args, "Oi", &arr, &node_capacity)) {
@@ -199,7 +172,7 @@ static PyObject* STRtree_new(PyTypeObject* type, PyObject* args, PyObject* kwds)
  * */
 
 void query_callback(void* item, void* user_data) {
-  kv_push(npy_intp, *(npy_intp_vec*)user_data, (npy_intp)item);
+  kv_push(npy_intp, *(index_vec_t*)user_data, (npy_intp)item);
 }
 
 /* Evaluate the predicate function against a prepared version of geom
@@ -233,8 +206,8 @@ void query_callback(void* item, void* user_data) {
 
 static char evaluate_predicate(void* context, FuncGEOS_YpY_b* predicate_func,
                                GEOSGeometry* geom, GEOSPreparedGeometry* prepared_geom,
-                               pg_geom_obj_vec* tree_geometries, npy_intp_vec* in_indexes,
-                               npy_intp_vec* out_indexes, npy_intp* count) {
+                               geom_obj_vec_t* tree_geometries, index_vec_t* in_indexes,
+                               index_vec_t* out_indexes, npy_intp* count) {
   GeometryObject* pg_geom;
   GEOSGeometry* target_geom;
   const GEOSPreparedGeometry* prepared_geom_tmp;
@@ -296,7 +269,7 @@ static PyObject* STRtree_query(STRtreeObject* self, PyObject* args) {
   int predicate_id = 0;  // default no predicate
   GEOSGeometry* geom = NULL;
   GEOSPreparedGeometry* prepared_geom = NULL;
-  npy_intp_vec query_indexes,
+  index_vec_t query_indexes,
       predicate_indexes;  // Resizable array for matches for each geometry
   npy_intp count;
   FuncGEOS_YpY_b* predicate_func = NULL;
@@ -340,7 +313,7 @@ static PyObject* STRtree_query(STRtreeObject* self, PyObject* args) {
   if (predicate_id == 0 || kv_size(query_indexes) == 0) {
     // No predicate function provided, return all geometry indexes from
     // query.  If array is empty, return an empty numpy array
-    result = copy_kvec_to_npy(&query_indexes);
+    result = index_vec_to_npy_arr(&query_indexes);
     kv_destroy(query_indexes);
     GEOS_FINISH;
     return (PyObject*)result;
@@ -357,7 +330,7 @@ static PyObject* STRtree_query(STRtreeObject* self, PyObject* args) {
     return NULL;
   }
 
-  result = copy_kvec_to_npy(&predicate_indexes);
+  result = index_vec_to_npy_arr(&predicate_indexes);
 
   kv_destroy(query_indexes);
   kv_destroy(predicate_indexes);
@@ -386,7 +359,7 @@ static PyObject* STRtree_query_bulk(STRtreeObject* self, PyObject* args) {
   int predicate_id = 0;  // default no predicate
   GEOSGeometry* geom = NULL;
   GEOSPreparedGeometry* prepared_geom = NULL;
-  npy_intp_vec query_indexes, src_indexes, target_indexes;
+  index_vec_t query_indexes, src_indexes, target_indexes;
   npy_intp i, j, n, size;
   FuncGEOS_YpY_b* predicate_func = NULL;
   PyArrayObject* result;

diff --git a/src/strtree.h b/src/strtree.h
@@ -5,23 +5,13 @@
 
 #include "geos.h"
 #include "pygeom.h"
+#include "vector.h"
 
-/* A resizable vector with numpy indices */
-typedef struct {
-  size_t n, m;
-  npy_intp* a;
-} npy_intp_vec;
-
-/* A resizable vector with pointers to pygeos GeometryObjects */
-typedef struct {
-  size_t n, m;
-  GeometryObject** a;
-} pg_geom_obj_vec;
 
 typedef struct {
   PyObject_HEAD void* ptr;
   npy_intp count;
-  pg_geom_obj_vec _geoms;
+  geom_obj_vec_t _geoms;
 } STRtreeObject;
 
 extern PyTypeObject STRtreeType;

diff --git a/src/vector.c b/src/vector.c
@@ -0,0 +1,43 @@
+#define PY_SSIZE_T_CLEAN
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+
+#include <Python.h>
+#include <structmember.h>
+
+#define NO_IMPORT_ARRAY
+#define PY_ARRAY_UNIQUE_SYMBOL pygeos_ARRAY_API
+
+#include <numpy/arrayobject.h>
+#include <numpy/ndarraytypes.h>
+#include <numpy/npy_3kcompat.h>
+
+#include "kvec.h"
+#include "vector.h"
+
+/* Copy values from arr of indexes to a new numpy integer array.
+ *
+ * Parameters
+ * ----------
+ * arr: dynamic vector array to convert to ndarray
+ */
+
+PyArrayObject* index_vec_to_npy_arr(index_vec_t* arr) {
+  Py_ssize_t i;
+  npy_intp size = kv_size(*arr);
+
+  npy_intp dims[1] = {size};
+  // the following raises a compiler warning based on how the macro is defined
+  // in numpy.  There doesn't appear to be anything we can do to avoid it.
+  PyArrayObject* result = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_INTP);
+  if (result == NULL) {
+    PyErr_SetString(PyExc_RuntimeError, "could not allocate numpy array");
+    return NULL;
+  }
+
+  for (i = 0; i < size; i++) {
+    // assign value into numpy array
+    *(npy_intp*)PyArray_GETPTR1(result, i) = kv_A(*arr, i);
+  }
+
+  return (PyArrayObject*)result;
+}
diff --git a/src/vector.h b/src/vector.h
@@ -0,0 +1,32 @@
+#ifndef _VECTOR_H
+#define _VECTOR_H
+
+#include <numpy/ndarraytypes.h>
+
+#include "pygeom.h"
+
+/* A resizable vector with numpy indices.
+ * Wraps the vector implementation in kvec.h as a type.
+ */
+typedef struct {
+  size_t n, m;
+  npy_intp* a;
+} index_vec_t;
+
+/* A resizable vector with pointers to pygeos GeometryObjects.
+ * Wraps the vector implementation in kvec.h as a type.
+ */
+typedef struct {
+  size_t n, m;
+  GeometryObject** a;
+} geom_obj_vec_t;
+
+/* Copy values from arr to a new numpy integer array.
+ *
+ * Parameters
+ * ----------
+ * arr: dynamic vector array to convert to ndarray
+ */
+extern PyArrayObject* index_vec_to_npy_arr(index_vec_t* arr);
+
+#endif