瀏覽代碼

putil: Assorted improvements to BitArray, SparseArray, *BitMask*:

* Support converting BitMask types to int
* BitArray constructor accepts a Python long of arbitrary size
* DoubleBitMask (and QuadBitMask, by extension) supports Python long in constructor
* Support for pickling (except DoubleBitMask)
* All of them now properly define __bool__()
* More unit tests

Fixes #886
rdb 5 年之前
父節點
當前提交
9d8c523dfa

+ 16 - 0
dtool/src/interrogatedb/py_compat.h

@@ -192,6 +192,22 @@ INLINE PyObject *_PyObject_FastCall(PyObject *func, PyObject **args, Py_ssize_t
   } while (0)
 #endif
 
+/* Python 3.8 */
+#if PY_VERSION_HEX < 0x03080000
+INLINE PyObject *_PyLong_Rshift(PyObject *a, size_t shiftby) {
+  PyObject *b = PyLong_FromLong(shiftby);
+  PyObject *result = PyNumber_Rshift(a, b);
+  Py_DECREF(b);
+  return result;
+}
+INLINE PyObject *_PyLong_Lshift(PyObject *a, size_t shiftby) {
+  PyObject *b = PyLong_FromLong(shiftby);
+  PyObject *result = PyNumber_Lshift(a, b);
+  Py_DECREF(b);
+  return result;
+}
+#endif
+
 /* Other Python implementations */
 
 // _PyErr_OCCURRED is an undocumented macro version of PyErr_Occurred.

+ 10 - 0
panda/src/putil/CMakeLists.txt

@@ -125,9 +125,19 @@ set(P3PUTIL_SOURCES
 set(P3PUTIL_IGATEEXT
   bamReader_ext.cxx
   bamReader_ext.h
+  bitArray_ext.cxx
+  bitArray_ext.h
+  bitArray_ext.I
+  bitMask_ext.h
+  bitMask_ext.I
   callbackObject_ext.h
+  doubleBitMask_ext.h
+  doubleBitMask_ext.I
   pythonCallbackObject.cxx
   pythonCallbackObject.h
+  sparseArray_ext.cxx
+  sparseArray_ext.h
+  sparseArray_ext.I
   typedWritable_ext.cxx
   typedWritable_ext.h
 )

+ 10 - 1
panda/src/putil/bitArray.h

@@ -21,6 +21,7 @@
 #include "typedObject.h"
 #include "indent.h"
 #include "pointerToArray.h"
+#include "extension.h"
 
 #include "checksumHashGenerator.h"
 
@@ -41,12 +42,14 @@ public:
   typedef BitMaskNative MaskType;
   typedef MaskType::WordType WordType;
 
+  INLINE BitArray(WordType init_value);
+
 PUBLISHED:
   enum { num_bits_per_word = MaskType::num_bits };
 
   INLINE BitArray();
-  INLINE BitArray(WordType init_value);
   BitArray(const SparseArray &from);
+  EXTENSION(BitArray(PyObject *init_value));
 
   INLINE static BitArray all_on();
   INLINE static BitArray all_off();
@@ -125,6 +128,10 @@ PUBLISHED:
   void operator <<= (int shift);
   void operator >>= (int shift);
 
+  EXTENSION(bool __bool__() const);
+  EXTENSION(PyObject *__getstate__() const);
+  EXTENSION(void __setstate__(PyObject *state));
+
 public:
   void generate_hash(ChecksumHashGenerator &hashgen) const;
 
@@ -138,6 +145,8 @@ private:
   Array _array;
   int _highest_bits;  // Either 0 or 1.
 
+  friend class Extension<BitArray>;
+
 public:
   void write_datagram(BamWriter *manager, Datagram &dg) const;
   void read_datagram(DatagramIterator &scan, BamReader *manager);

+ 20 - 0
panda/src/putil/bitArray_ext.I

@@ -0,0 +1,20 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file bitArray_ext.I
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+/**
+ * Returns true if the value is not zero.
+ */
+INLINE bool Extension<BitArray>::
+__bool__() const {
+  return !_this->is_zero();
+}

+ 101 - 0
panda/src/putil/bitArray_ext.cxx

@@ -0,0 +1,101 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file bitArray_ext.cxx
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+#include "bitArray_ext.h"
+
+#ifdef HAVE_PYTHON
+
+/**
+ * Creates a BitArray from a Python long object.
+ */
+void Extension<BitArray>::
+__init__(PyObject *init_value) {
+#if PY_MAJOR_VERSION < 3
+  if (PyInt_Check(init_value)) {
+    long value = PyInt_AS_LONG(init_value);
+    if (value >= 0) {
+      _this->set_word(0, value);
+    } else {
+      PyErr_SetString(PyExc_ValueError, "BitArray constructor requires a positive integer");
+    }
+    return;
+  }
+#endif
+
+  if (!PyLong_Check(init_value) || Py_SIZE(init_value) < 0) {
+    PyErr_SetString(PyExc_ValueError, "BitArray constructor requires a positive integer");
+    return;
+  }
+
+  int n = _PyLong_NumBits(init_value);
+  if (n > 0) {
+    size_t num_words = (n + BitArray::num_bits_per_word - 1) / BitArray::num_bits_per_word;
+    _this->_array.resize(num_words);
+    _PyLong_AsByteArray((PyLongObject *)init_value,
+      (unsigned char *)&_this->_array[0],
+      num_words * sizeof(BitArray::WordType),
+      1, 0);
+  }
+}
+
+/**
+ * Returns the value of the BitArray as a picklable Python object.
+ *
+ * We could just return a list of words.  However, different builds of Panda3D
+ * may have different sizes for the WordType, so we'd also need to add code to
+ * convert between different WordTypes.  Instead, we'll just encode the whole
+ * array as a Python long, with infinite arrays stored as inverted longs.
+ */
+PyObject *Extension<BitArray>::
+__getstate__() const {
+  if (_this->_array.empty()) {
+    return PyLong_FromLong(-_this->_highest_bits);
+  }
+
+  if (_this->_highest_bits == 0) {
+    return _PyLong_FromByteArray(
+      (const unsigned char *)&_this->_array[0],
+      _this->_array.size() * sizeof(BitArray::WordType),
+      1, 0);
+  } else {
+    // This is an infinite array, so we invert it to make it a finite array and
+    // store it as an inverted long.
+    BitArray copy(*_this);
+    copy.invert_in_place();
+    PyObject *state = _PyLong_FromByteArray(
+      (const unsigned char *)&copy._array[0],
+      copy._array.size() * sizeof(BitArray::WordType),
+      1, 0);
+    PyObject *inverted = PyNumber_Invert(state);
+    Py_DECREF(state);
+    return inverted;
+  }
+}
+
+/**
+ * Takes the value returned by __getstate__ and uses it to freshly initialize
+ * this BitArray object.
+ */
+void Extension<BitArray>::
+__setstate__(PyObject *state) {
+  if (Py_SIZE(state) >= 0) {
+    __init__(state);
+  } else {
+    PyObject *inverted = PyNumber_Invert(state);
+    __init__(inverted);
+    Py_DECREF(inverted);
+    _this->invert_in_place();
+  }
+}
+
+#endif

+ 44 - 0
panda/src/putil/bitArray_ext.h

@@ -0,0 +1,44 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file bitArray_ext.h
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+#ifndef BITARRAY_EXT_H
+#define BITARRAY_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "bitArray.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for BitArray, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<>
+class Extension<BitArray> : public ExtensionBase<BitArray> {
+public:
+  void __init__(PyObject *init_value);
+
+  INLINE bool __bool__() const;
+
+  PyObject *__getstate__() const;
+  void __setstate__(PyObject *state);
+};
+
+#include "bitArray_ext.I"
+
+#endif  // HAVE_PYTHON
+
+#endif  // BITARRAY_EXT_H

+ 0 - 9
panda/src/putil/bitMask.I

@@ -623,15 +623,6 @@ get_key() const {
   return (int)_word;
 }
 
-/**
- * Returns true if the bitmask is not zero.
- */
-template<class WType, int nbits>
-INLINE bool BitMask<WType, nbits>::
-__nonzero__() const {
-  return _word != 0;
-}
-
 /**
  * Adds the bitmask to the indicated hash generator.
  */

+ 3 - 1
panda/src/putil/bitMask.h

@@ -125,7 +125,9 @@ PUBLISHED:
 
   INLINE int get_key() const;
 
-  INLINE bool __nonzero__() const;
+  EXTENSION(bool __bool__() const);
+  EXTENSION(PyObject *__int__() const);
+  EXTENSION(PyObject *__reduce__(PyObject *self) const);
 
 public:
   INLINE void generate_hash(ChecksumHashGenerator &hashgen) const;

+ 43 - 0
panda/src/putil/bitMask_ext.I

@@ -0,0 +1,43 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file bitMask_ext.I
+ * @author rdb
+ * @date 2020-03-22
+ */
+
+/**
+ * Returns true if the value is not zero.
+ */
+template<class WType, int nbits>
+INLINE bool Extension<BitMask<WType, nbits> >::
+__bool__() const {
+  return this->_this->get_word() != 0;
+}
+
+/**
+ * Returns the value as an integer.
+ */
+template<class WType, int nbits>
+INLINE PyObject *Extension<BitMask<WType, nbits> >::
+__int__() const {
+  return Dtool_WrapValue(this->_this->get_word());
+}
+
+/**
+ * This special Python method is implemented to provide support for the pickle
+ * module.
+ */
+template<class WType, int nbits>
+INLINE PyObject *Extension<BitMask<WType, nbits> >::
+__reduce__(PyObject *self) const {
+  // We should return at least a 2-tuple, (Class, (args)): the necessary class
+  // object whose constructor we should call (e.g.  this), and the arguments
+  // necessary to reconstruct this object.
+  return Py_BuildValue("(O(k))", Py_TYPE(self), this->_this->get_word());
+}

+ 41 - 0
panda/src/putil/bitMask_ext.h

@@ -0,0 +1,41 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file bitMask_ext.h
+ * @author rdb
+ * @date 2020-03-22
+ */
+
+#ifndef BITMASK_EXT_H
+#define BITMASK_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "bitMask.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for BitMask, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<class WType, int nbits>
+class Extension<BitMask<WType, nbits> > : public ExtensionBase<BitMask<WType, nbits> > {
+public:
+  INLINE bool __bool__() const;
+  INLINE PyObject *__int__() const;
+  INLINE PyObject *__reduce__(PyObject *self) const;
+};
+
+#include "bitMask_ext.I"
+
+#endif  // HAVE_PYTHON
+
+#endif  // BITMASK_EXT_H

+ 8 - 0
panda/src/putil/doubleBitMask.h

@@ -17,6 +17,7 @@
 #include "pandabase.h"
 
 #include "bitMask.h"
+#include "extension.h"
 
 /**
  * This is a special BitMask type that is implemented as a pair of lesser
@@ -38,6 +39,7 @@ PUBLISHED:
   };
 
   constexpr DoubleBitMask() = default;
+  EXTENSION(DoubleBitMask(PyObject *init_value));
 
   INLINE static DoubleBitMask<BMType> all_on();
   INLINE static DoubleBitMask<BMType> all_off();
@@ -110,12 +112,18 @@ PUBLISHED:
   INLINE void operator <<= (int shift);
   INLINE void operator >>= (int shift);
 
+  EXTENSION(bool __bool__() const);
+  EXTENSION(PyObject *__int__() const);
+  EXTENSION(PyObject *__reduce__(PyObject *self) const);
+
 public:
   INLINE void generate_hash(ChecksumHashGenerator &hashgen) const;
 
 private:
   BitMaskType _lo, _hi;
 
+  friend class Extension<DoubleBitMask>;
+
 public:
   static TypeHandle get_class_type() {
     return _type_handle;

+ 93 - 0
panda/src/putil/doubleBitMask_ext.I

@@ -0,0 +1,93 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file doubleBitMask_ext.I
+ * @author rdb
+ * @date 2020-04-01
+ */
+
+/**
+ * Initializes a DoubleBitMask from a Python long integer.
+ */
+template<class BMType>
+INLINE void Extension<DoubleBitMask<BMType> >::
+__init__(PyObject *init_value) {
+#if PY_MAJOR_VERSION < 3
+  if (PyInt_Check(init_value)) {
+    long value = PyInt_AS_LONG(init_value);
+    if (value >= 0) {
+      this->_this->store((typename BMType::WordType)value, 0, sizeof(long) * 8 - 1);
+    } else {
+      PyErr_SetString(PyExc_ValueError, "DoubleBitMask constructor requires a positive integer");
+    }
+    return;
+  }
+#endif
+
+  if (!PyLong_Check(init_value) || Py_SIZE(init_value) < 0) {
+    PyErr_SetString(PyExc_ValueError, "DoubleBitMask constructor requires a positive integer");
+    return;
+  }
+
+  int n = _PyLong_NumBits(init_value);
+  if (n > DoubleBitMask<BMType>::num_bits) {
+    PyErr_SetString(PyExc_OverflowError, "value out of range for DoubleBitMask");
+    return;
+  }
+
+  if (n > 0) {
+    size_t num_bytes = (n + 7) / 8;
+    unsigned char *bytes = (unsigned char *)alloca(num_bytes);
+    _PyLong_AsByteArray((PyLongObject *)init_value, bytes, num_bytes, 1, 0);
+
+    for (size_t i = 0; i < num_bytes; ++i) {
+      this->_this->store(bytes[i], i * 8, 8);
+    }
+  }
+}
+
+/**
+ * Returns true if the value is not zero.
+ */
+template<class BMType>
+INLINE bool Extension<DoubleBitMask<BMType> >::
+__bool__() const {
+  return !this->_this->is_zero();
+}
+
+/**
+ * Returns the value as an integer.
+ */
+template<class BMType>
+INLINE PyObject *Extension<DoubleBitMask<BMType> >::
+__int__() const {
+  PyObject *result = invoke_extension(&this->_this->_lo).__int__();
+  if (!this->_this->_hi.is_zero()) {
+    PyObject *lo = result;
+    PyObject *hi = invoke_extension(&this->_this->_hi).__int__();
+    PyObject *shifted = _PyLong_Lshift(hi, DoubleBitMask<BMType>::half_bits);
+    Py_DECREF(hi);
+    result = PyNumber_Or(shifted, lo);
+    Py_DECREF(shifted);
+    Py_DECREF(lo);
+  }
+  return result;
+}
+
+/**
+ * This special Python method is implemented to provide support for the pickle
+ * module.
+ */
+template<class BMType>
+INLINE PyObject *Extension<DoubleBitMask<BMType> >::
+__reduce__(PyObject *self) const {
+  // We should return at least a 2-tuple, (Class, (args)): the necessary class
+  // object whose constructor we should call (e.g.  this), and the arguments
+  // necessary to reconstruct this object.
+  return Py_BuildValue("(O(N))", Py_TYPE(self), __int__());
+}

+ 45 - 0
panda/src/putil/doubleBitMask_ext.h

@@ -0,0 +1,45 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file doubleBitMask_ext.h
+ * @author rdb
+ * @date 2020-04-01
+ */
+
+#ifndef DOUBLEBITMASK_EXT_H
+#define DOUBLEBITMASK_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "doubleBitMask.h"
+#include "py_panda.h"
+
+#include "bitMask_ext.h"
+
+/**
+ * This class defines the extension methods for DoubleBitMask, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<class BMType>
+class Extension<DoubleBitMask<BMType> > : public ExtensionBase<DoubleBitMask<BMType> > {
+public:
+  INLINE void __init__(PyObject *init_value);
+
+  INLINE bool __bool__() const;
+  INLINE PyObject *__int__() const;
+  INLINE PyObject *__reduce__(PyObject *self) const;
+};
+
+#include "doubleBitMask_ext.I"
+
+#endif  // HAVE_PYTHON
+
+#endif  // DOUBLEBITMASK_EXT_H

+ 2 - 0
panda/src/putil/p3putil_ext_composite.cxx

@@ -1,3 +1,5 @@
 #include "bamReader_ext.cxx"
+#include "bitArray_ext.cxx"
 #include "pythonCallbackObject.cxx"
+#include "sparseArray_ext.cxx"
 #include "typedWritable_ext.cxx"

+ 7 - 0
panda/src/putil/sparseArray.h

@@ -16,6 +16,7 @@
 
 #include "pandabase.h"
 #include "ordered_vector.h"
+#include "extension.h"
 
 class BitArray;
 class BamWriter;
@@ -116,6 +117,10 @@ PUBLISHED:
   INLINE int get_subrange_begin(size_t n) const;
   INLINE int get_subrange_end(size_t n) const;
 
+  EXTENSION(bool __bool__() const);
+  EXTENSION(PyObject *__getstate__() const);
+  EXTENSION(void __setstate__(PyObject *state));
+
 private:
   void do_add_range(int begin, int end);
   void do_remove_range(int begin, int end);
@@ -140,6 +145,8 @@ private:
   Subranges _subranges;
   bool _inverse;
 
+  friend class Extension<SparseArray>;
+
 public:
   void write_datagram(BamWriter *manager, Datagram &dg) const;
   void read_datagram(DatagramIterator &scan, BamReader *manager);

+ 20 - 0
panda/src/putil/sparseArray_ext.I

@@ -0,0 +1,20 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file sparseArray_ext.I
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+/**
+ * Returns true if the value is not zero.
+ */
+INLINE bool Extension<SparseArray>::
+__bool__() const {
+  return !_this->is_zero();
+}

+ 93 - 0
panda/src/putil/sparseArray_ext.cxx

@@ -0,0 +1,93 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file sparseArray_ext.cxx
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+#include "sparseArray_ext.h"
+
+#ifdef HAVE_PYTHON
+
+/**
+ * Returns the value of the SparseArray as a picklable Python object.
+ *
+ * We store this as a tuple of integers.  The first number indicates the first
+ * bit that is set to 1, the second number indicates the next bit that is set to
+ * 0, the third number indicates the next bit that is set to 1, etc.  Using this
+ * method, we store an uneven number of integers for an inverted SparseArray,
+ * and an even number for a regular SparseArray.
+ *
+ * This table demonstrates the three different cases:
+ *
+ *  | range      | pickled     | SparseArray |
+ *  |------------|-------------|-------------|
+ *  | 0-2, 4-8   | 0, 2, 4, 8  | 0-2, 4-8    |
+ *  | 0-2, 4-... | 0, 2, 4     | ~ 2-4       |
+ *  | 2-4, 6-... | 2, 4, 6     | ~ 0-2, 4-6  |
+ *
+ */
+PyObject *Extension<SparseArray>::
+__getstate__() const {
+  PyObject *state;
+  Py_ssize_t index = 0;
+  size_t sri = 0;
+  size_t num_ranges = _this->get_num_subranges();
+
+  if (!_this->is_inverse()) {
+    state = PyTuple_New(num_ranges * 2);
+  }
+  else if (num_ranges > 0 && _this->get_subrange_begin(0) == 0) {
+    // Prevent adding a useless 0-0 range at the beginning.
+    state = PyTuple_New(num_ranges * 2 - 1);
+    PyTuple_SET_ITEM(state, index++, Dtool_WrapValue(_this->get_subrange_end(sri++)));
+  }
+  else {
+    state = PyTuple_New(num_ranges * 2 + 1);
+    PyTuple_SET_ITEM(state, index++, Dtool_WrapValue(0));
+  }
+
+  for (; sri < num_ranges; ++sri) {
+    PyTuple_SET_ITEM(state, index++, Dtool_WrapValue(_this->get_subrange_begin(sri)));
+    PyTuple_SET_ITEM(state, index++, Dtool_WrapValue(_this->get_subrange_end(sri)));
+  }
+  return state;
+}
+
+/**
+ * Takes the tuple returned by __getstate__ and uses it to freshly initialize
+ * this SparseArray object.
+ */
+void Extension<SparseArray>::
+__setstate__(PyObject *state) {
+  _this->clear();
+
+  Py_ssize_t i = 0;
+  Py_ssize_t len = PyTuple_GET_SIZE(state);
+  if (len % 2 != 0) {
+    // An uneven number of elements indicates an open final range.
+    // This translates to an inverted range in SparseArray's representation.
+    _this->invert_in_place();
+    long first = PyLongOrInt_AS_LONG(PyTuple_GET_ITEM(state, 0));
+    if (first != 0) {
+      // It doesn't start at 0, so we have to first disable this range.
+      _this->do_add_range(0, (int)first);
+    }
+    ++i;
+  }
+
+  for (; i < len; i += 2) {
+    _this->do_add_range(
+      PyLongOrInt_AS_LONG(PyTuple_GET_ITEM(state, i)),
+      PyLongOrInt_AS_LONG(PyTuple_GET_ITEM(state, i + 1))
+    );
+  }
+}
+
+#endif

+ 42 - 0
panda/src/putil/sparseArray_ext.h

@@ -0,0 +1,42 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file sparseArray_ext.h
+ * @author rdb
+ * @date 2020-03-21
+ */
+
+#ifndef SPARSEARRAY_EXT_H
+#define SPARSEARRAY_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "sparseArray.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for SparseArray, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<>
+class Extension<SparseArray> : public ExtensionBase<SparseArray> {
+public:
+  INLINE bool __bool__() const;
+
+  PyObject *__getstate__() const;
+  void __setstate__(PyObject *state);
+};
+
+#include "sparseArray_ext.I"
+
+#endif  // HAVE_PYTHON
+
+#endif  // SPARSEARRAY_EXT_H

+ 60 - 0
tests/putil/test_bitarray.py

@@ -0,0 +1,60 @@
+from panda3d.core import BitArray
+import pickle
+import pytest
+
+
+def test_bitarray_constructor():
+    assert BitArray().is_zero()
+    assert BitArray(0).is_zero()
+
+    ba = BitArray(0x10000000000000000000000000)
+    assert ba.get_lowest_on_bit() == 100
+    assert ba.get_highest_on_bit() == 100
+
+    with pytest.raises(Exception):
+        assert BitArray(-1)
+
+    with pytest.raises(Exception):
+        assert BitArray(-10000000000000000000)
+
+
+def test_bitarray_allon():
+    assert BitArray.all_on().is_all_on()
+
+
+def test_bitarray_nonzero():
+    assert not BitArray()
+    assert not BitArray(0)
+    assert BitArray(1)
+    assert BitArray.all_on()
+
+
+def test_bitarray_invert():
+    assert ~BitArray(0) != BitArray(0)
+    assert (~BitArray(0)).is_all_on()
+    assert ~~BitArray(0) == BitArray(0)
+    assert ~~BitArray(123) == BitArray(123)
+
+
+def test_bitarray_getstate():
+    assert BitArray().__getstate__() == 0
+    assert BitArray(0).__getstate__() == 0
+    assert BitArray(100).__getstate__() == 100
+    assert BitArray(9870000000000000000).__getstate__() == 9870000000000000000
+    assert BitArray.all_on().__getstate__() == -1
+    assert ~BitArray(100).__getstate__() == ~100
+    assert ~BitArray(812000000000000000).__getstate__() == ~812000000000000000
+
+
+def test_bitarray_pickle():
+    ba = BitArray()
+    assert ba == pickle.loads(pickle.dumps(ba, -1))
+
+    ba = BitArray(0)
+    assert ba == pickle.loads(pickle.dumps(ba, -1))
+
+    ba = BitArray(123)
+    assert ba == pickle.loads(pickle.dumps(ba, -1))
+
+    ba = BitArray(94187049178237918273981729127381723)
+    assert ba == pickle.loads(pickle.dumps(ba, -1))

+ 67 - 7
tests/putil/test_bitmask.py

@@ -1,10 +1,70 @@
-from panda3d import core
+from panda3d.core import BitMask16, BitMask32, BitMask64
+from panda3d.core import DoubleBitMaskNative, QuadBitMaskNative
+import pickle
+import pytest
+
+
+double_num_bits = DoubleBitMaskNative.get_max_num_bits()
+quad_num_bits = QuadBitMaskNative.get_max_num_bits()
 
 
 def test_bitmask_allon():
-    assert core.BitMask16.all_on().is_all_on()
-    assert core.BitMask32.all_on().is_all_on()
-    assert core.BitMask64.all_on().is_all_on()
-    assert core.DoubleBitMaskNative.all_on().is_all_on()
-    assert core.QuadBitMaskNative.all_on().is_all_on()
-    assert core.BitArray.all_on().is_all_on()
+    assert BitMask16.all_on().is_all_on()
+    assert BitMask32.all_on().is_all_on()
+    assert BitMask64.all_on().is_all_on()
+    assert DoubleBitMaskNative.all_on().is_all_on()
+    assert QuadBitMaskNative.all_on().is_all_on()
+
+    assert DoubleBitMaskNative((1 << double_num_bits) - 1).is_all_on()
+    assert QuadBitMaskNative((1 << quad_num_bits) - 1).is_all_on()
+
+
+def test_bitmask_nonzero():
+    assert not BitMask16()
+    assert not BitMask32()
+    assert not BitMask64()
+    assert not DoubleBitMaskNative()
+    assert not QuadBitMaskNative()
+
+
+def test_bitmask_overflow():
+    with pytest.raises(OverflowError):
+        DoubleBitMaskNative(1 << double_num_bits)
+
+    with pytest.raises(OverflowError):
+        QuadBitMaskNative(1 << quad_num_bits)
+
+
+def test_bitmask_int():
+    assert int(BitMask16()) == 0
+    assert int(BitMask16(0xfffe)) == 0xfffe
+
+    assert int(BitMask32()) == 0
+    assert int(BitMask32(1)) == 1
+    assert int(BitMask32(1234567)) == 1234567
+    assert int(BitMask32(0x8ff123fe)) == 0x8ff123fe
+    assert int(BitMask32(0xffffffff)) == 0xffffffff
+
+    assert int(BitMask64()) == 0
+    assert int(BitMask64(0xfffffffffffffffe)) == 0xfffffffffffffffe
+
+    assert int(DoubleBitMaskNative()) == 0
+    assert int(DoubleBitMaskNative(1)) == 1
+    assert int(DoubleBitMaskNative(1 << (double_num_bits // 2))) == 1 << (double_num_bits // 2)
+    assert int(DoubleBitMaskNative(0xffff0001)) == 0xffff0001
+    assert int(DoubleBitMaskNative((1 << double_num_bits) - 1)) == (1 << double_num_bits) - 1
+
+    assert int(QuadBitMaskNative()) == 0
+    assert int(QuadBitMaskNative(1)) == 1
+    assert int(QuadBitMaskNative(1 << (quad_num_bits // 2))) == 1 << (quad_num_bits // 2)
+    assert int(QuadBitMaskNative(0xffff0001feff0002)) == 0xffff0001feff0002
+    assert int(QuadBitMaskNative((1 << quad_num_bits) - 1)) == (1 << quad_num_bits) - 1
+
+
+def test_bitmask_pickle():
+    assert pickle.loads(pickle.dumps(BitMask16(0), -1)).is_zero()
+
+    mask1 = BitMask16(123)
+    data = pickle.dumps(mask1, -1)
+    mask2 = pickle.loads(data)
+    assert mask1 == mask2

+ 63 - 0
tests/putil/test_sparsearray.py

@@ -1,4 +1,5 @@
 from panda3d import core
+import pickle
 
 
 def test_sparse_array_set_bit_to():
@@ -232,3 +233,65 @@ def test_sparse_array_augm_assignment():
     u = core.SparseArray()
     t ^= u
     assert s is t
+
+
+def test_sparse_array_nonzero():
+    sa = core.SparseArray()
+    assert not sa
+    sa.set_bit(0)
+    assert sa
+
+    sa = core.SparseArray.all_on()
+    assert sa
+    sa.clear_range(0, 100)
+    assert sa
+
+
+def test_sparse_array_getstate():
+    sa = core.SparseArray()
+    assert sa.__getstate__() == ()
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    assert sa.__getstate__() == (0,)
+
+    sa = core.SparseArray()
+    sa.set_range(0, 2)
+    sa.set_range(4, 4)
+    assert sa.__getstate__() == (0, 2, 4, 8)
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    sa.clear_range(2, 4)
+    assert sa.__getstate__() == (0, 2, 6)
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    sa.clear_range(0, 2)
+    sa.clear_range(4, 4)
+    assert sa.__getstate__() == (2, 4, 8)
+
+
+def test_sparse_array_pickle():
+    sa = core.SparseArray()
+    assert sa == pickle.loads(pickle.dumps(sa, -1))
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    assert sa == pickle.loads(pickle.dumps(sa, -1))
+
+    sa = core.SparseArray()
+    sa.set_range(0, 2)
+    sa.set_range(4, 4)
+    assert sa == pickle.loads(pickle.dumps(sa, -1))
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    sa.clear_range(2, 4)
+    assert sa == pickle.loads(pickle.dumps(sa, -1))
+
+    sa = core.SparseArray()
+    sa.invert_in_place()
+    sa.clear_range(0, 2)
+    sa.clear_range(4, 4)
+    assert sa == pickle.loads(pickle.dumps(sa, -1))