add step argument to count method

bitarray/_bitarray.c

@@ -948,26 +948,73 @@ PyDoc_STRVAR(copy_doc,
 Return a copy of the bitarray.");
 
 
+static Py_ssize_t
+calc_slicelength(Py_ssize_t start, Py_ssize_t stop, Py_ssize_t step)
+{
+ assert(step != 0);
+ if (step < 0) {
+ if (stop < start)
+ return (start - stop - 1) / (-step) + 1;
+ }
+ else {
+ if (start < stop)
+ return (stop - start - 1) / step + 1;
+ }
+ return 0;
+}
+
+/* adjust slice parameters such that step is always positive */
+static void
+make_step_positive(Py_ssize_t *start, Py_ssize_t *stop, Py_ssize_t *step)
+{
+ if (*step < 0) {
+ Py_ssize_t slicelength = calc_slicelength(*start, *stop, *step);
+
+ *stop = *start + 1;
+ *start = *stop + *step * (slicelength - 1) - 1;
+ *step *= -1;
+ assert(*start < *stop || slicelength == 0);
+ }
+ assert(*start >= 0 && *stop >= 0);
+}
+
 static PyObject *
 bitarray_count(bitarrayobject *self, PyObject *args)
 {
  PyObject *value = Py_True;
- Py_ssize_t start = 0, stop = self->nbits;
+ Py_ssize_t start = 0, stop = self->nbits, step = 1;
  int vi;
 
- if (!PyArg_ParseTuple(args, "|Onn:count", &value, &start, &stop))
+ if (!PyArg_ParseTuple(args, "|Onnn:count", &value, &start, &stop, &step))
  return NULL;
  if ((vi = pybit_as_int(value)) < 0)
  return NULL;
 
- normalize_index(self->nbits, &start);
- normalize_index(self->nbits, &stop);
+ normalize_index(self->nbits, step, &start);
+ normalize_index(self->nbits, step, &stop);
+
+ if (step == 1) {
+ return PyLong_FromSsize_t(count(self, vi, start, stop));
+ }
+ else if (step == 0) {
+ PyErr_SetString(PyExc_ValueError, "count step cannot be zero");
+ return NULL;
+ }
+ else {
+ Py_ssize_t cnt = 0, i;
+
+ make_step_positive(&start, &stop, &step);
+ for (i = start; i < stop; i += step)
+ cnt += getbit((bitarrayobject *) self, i);
 
- return PyLong_FromSsize_t(count(self, vi, start, stop));
+ if (!vi)
+ cnt = calc_slicelength(start, stop, step) - cnt;
+ return PyLong_FromSsize_t(cnt);
+ }
 }
 
 PyDoc_STRVAR(count_doc,
-"count(value=1, start=0, stop=<end of array>, /) -> int\n\
+"count(value=1, start=0, stop=<end of array>, step=1, /) -> int\n\
 \n\
 Count the number of occurrences of `value` in the bitarray.");
 
@@ -1032,8 +1079,8 @@ bitarray_find(bitarrayobject *self, PyObject *args)
  if (!PyArg_ParseTuple(args, "O|nn", &x, &start, &stop))
  return NULL;
 
- normalize_index(self->nbits, &start);
- normalize_index(self->nbits, &stop);
+ normalize_index(self->nbits, 1, &start);
+ normalize_index(self->nbits, 1, &stop);
 
  if (PyIndex_Check(x)) {
  int vi;
@@ -1101,7 +1148,7 @@ bitarray_insert(bitarrayobject *self, PyObject *args)
  if (!PyArg_ParseTuple(args, "nO:insert", &i, &value))
  return NULL;
 
- normalize_index(self->nbits, &i);
+ normalize_index(self->nbits, 1, &i);
 
  vi = pybit_as_int(value);
  if (vi < 0)
@@ -1974,7 +2021,7 @@ setslice_bitarray(bitarrayobject *self, PyObject *slice,
  return res;
 }
 
-/* like PySlice_GetIndicesEx(), but step index will always be positive */
+/* like PySlice_GetIndicesEx(), but step will always be positive */
 static int
 slice_get_indices(PyObject *slice, Py_ssize_t length,
  Py_ssize_t *start, Py_ssize_t *stop, Py_ssize_t *step,
@@ -1985,19 +2032,8 @@ slice_get_indices(PyObject *slice, Py_ssize_t length,
  start, stop, step, slicelength) < 0)
  return -1;
 
- if (*slicelength == 0)
- return 0;
-
- if (*step < 0) {
- *stop = *start + 1;
- *start = *stop + *step * (*slicelength - 1) - 1;
- *step *= -1;
- }
- assert(*step > 0 && *start < *stop && *slicelength > 0);
- assert(0 <= *start && *start < length);
- assert(0 <= *stop && *stop <= length);
- assert(*step != 1 || *start + *slicelength == *stop);
- assert(*start + ((*slicelength - 1) * *step) < *stop);
+ assert(calc_slicelength(*start, *stop, *step) == *slicelength);
+ make_step_positive(start, stop, step);
  return 0;
 }
 

bitarray/_util.c

@@ -215,8 +215,8 @@ r_index(PyObject *module, PyObject *args)
  return NULL;
 
 #define aa ((bitarrayobject *) a)
- normalize_index(aa->nbits, &start);
- normalize_index(aa->nbits, &stop);
+ normalize_index(aa->nbits, 1, &start);
+ normalize_index(aa->nbits, 1, &stop);
  res = find_last(aa, vi, start, stop);
 #undef aa
  if (res < 0)

bitarray/bitarray.h

@@ -154,15 +154,16 @@ static const unsigned char bitcount_lookup[256] = {
 
 /* normalize index (which may be negative), such that 0 <= i <= n */
 static inline void
-normalize_index(Py_ssize_t n, Py_ssize_t *i)
+normalize_index(Py_ssize_t length, Py_ssize_t step, Py_ssize_t *i)
 {
  if (*i < 0) {
- *i += n;
+ *i += length;
  if (*i < 0)
- *i = 0;
+ *i = (step < 0) ? -1 : 0;
+ }
+ else if (*i >= length) {
+ *i = (step < 0) ? length - 1 : length;
  }
- if (*i > n)
- *i = n;
 }
 
 /* Interpret a PyObject (usually PyLong or PyBool) as a bit, return 0 or 1.

bitarray/test_bitarray.py

@@ -2711,6 +2711,7 @@ def test_basic(self):
  self.assertEqual(a.count(1), 3)
  self.assertEqual(a.count(0), 2)
  self.assertRaises(ValueError, a.count, 2)
+ self.assertRaises(ValueError, a.count, 1, 0, 5, 0)
  self.assertRaises(TypeError, a.count, None)
  self.assertRaises(TypeError, a.count, '')
  self.assertRaises(TypeError, a.count, 'A')
@@ -2728,9 +2729,12 @@ def test_byte(self):
 
  def test_whole_range(self):
  for a in self.randombitarrays():
+ n = len(a)
  s = a.to01()
- self.assertEqual(a.count(1), s.count('1'))
- self.assertEqual(a.count(0), s.count('0'))
+ for v in 0, 1:
+ ref = s.count(str(v))
+ self.assertEqual(a.count(v), ref)
+ self.assertEqual(a.count(v, n, -n - 1, -1), ref)
 
  def test_zeros(self):
  N = 37
@@ -2743,6 +2747,9 @@ def test_explicit(self):
  a = bitarray('01001100 01110011 01')
  self.assertEqual(a.count(), 9)
  self.assertEqual(a.count(0, 12), 3)
+ self.assertEqual(a.count(1, 1, 18, 2), 6)
+ self.assertEqual(a.count(1, 0, 18, 3), 2)
+ self.assertEqual(a.count(1, 15, 4, -3), 2)
  self.assertEqual(a.count(1, -5), 3)
  self.assertEqual(a.count(1, 2, 17), 7)
  self.assertEqual(a.count(1, 6, 11), 2)
@@ -2752,11 +2759,11 @@ def test_explicit(self):
 
  def test_random(self):
  for a in self.randombitarrays():
- s = a.to01()
  i = randint(-3, len(a) + 2)
  j = randint(-3, len(a) + 2)
- self.assertEqual(a.count(1, i, j), s[i:j].count('1'))
- self.assertEqual(a.count(0, i, j), s[i:j].count('0'))
+ for s in 1, 2, randint(1, len(a) + 1):
+ for v in 0, 1:
+ self.assertEqual(a.count(v, i, j, s), a[i:j:s].count(v))
 
 tests.append(CountTests)
 

update_doc.py

@@ -15,7 +15,8 @@
  'bitarray': '2.3: optional `buffer` argument',
  'bitarray.bytereverse': '2.2.5: optional start and stop arguments',
  'bitarray.clear': '1.4',
- 'bitarray.count': '1.1.0: optional start and stop arguments',
+ 'bitarray.count': ['1.1.0: optional start and stop arguments',
+ '2.3.7: optional step argument'],
  'bitarray.find': '2.1',
  'bitarray.invert': '1.5.3: optional index argument',
  'decodetree': '1.6',
@@ -77,7 +78,8 @@ def write_doc(fo, name):
 
  new_in = NEW_IN.get(name)
  if new_in:
- fo.write("\n New in version %s.\n" % new_in.replace('`', '``'))
+ for line in new_in if isinstance(new_in, list) else [new_in]:
+ fo.write("\n New in version %s.\n" % line.replace('`', '``'))
 
  fo.write('\n\n')