diff --git a/dynamicsizerecarray/__init__.py b/dynamicsizerecarray/__init__.py index 6affc73..919fc50 100644 --- a/dynamicsizerecarray/__init__.py +++ b/dynamicsizerecarray/__init__.py @@ -8,7 +8,7 @@ class DynamicSizeRecarray: A dynamic, appendable implementation of numpy.core.records.recarray. """ - def __init__(self, recarray=None, dtype=None): + def __init__(self, recarray=None, dtype=None, shape=0): """ Either provide an existing recarray 'recarray' or provide the 'dtype' to start with an empty recarray. @@ -19,6 +19,8 @@ def __init__(self, recarray=None, dtype=None): The start of the dynamic recarray. dtype : list(tuple("key", "dtype_str")), default=None The dtype of the dynamic recarray. + shape : int + The initial size. (Same as in np.core.records.recarray) """ self._size = 0 if recarray is None and dtype == None: @@ -28,19 +30,24 @@ def __init__(self, recarray=None, dtype=None): "Expected either one of 'recarray' or' dtype' to be 'None'" ) + _minimal_capacity = 2 + if dtype: - recarray = np.core.records.recarray( - shape=0, + if shape < 0: + raise AttributeError("Expected shape >= 0.") + initial_capacity = np.max([_minimal_capacity, shape]) + self._recarray = np.core.records.recarray( + shape=initial_capacity, dtype=dtype, ) - - _minimal_capacity = 2 - initial_capacity = np.max([_minimal_capacity, len(recarray)]) - self._recarray = np.core.records.recarray( - shape=initial_capacity, - dtype=recarray.dtype, - ) - self.append_recarray(recarray=recarray) + self._size = shape + else: + initial_capacity = np.max([_minimal_capacity, len(recarray)]) + self._recarray = np.core.records.recarray( + shape=initial_capacity, + dtype=recarray.dtype, + ) + self.append_recarray(recarray=recarray) @property def shape(self): diff --git a/dynamicsizerecarray/tests/test_init_shape.py b/dynamicsizerecarray/tests/test_init_shape.py new file mode 100644 index 0000000..9233cc1 --- /dev/null +++ b/dynamicsizerecarray/tests/test_init_shape.py @@ -0,0 +1,21 @@ +import dynamicsizerecarray +import pytest +import numpy as np + +def test_init_with_shape(): + dra = dynamicsizerecarray.DynamicSizeRecarray( + dtype=[("a", "i8"), ("b", "u2")], + shape=13, + ) + + assert len(dra) == 13 + assert dra._capacity() == 13 + + +def test_init_negative_shape(): + with pytest.raises(AttributeError): + dra = dynamicsizerecarray.DynamicSizeRecarray( + dtype=[("a", "i8")], + shape=-1, + ) + diff --git a/dynamicsizerecarray/version.py b/dynamicsizerecarray/version.py index 00ec2dc..9b36b86 100644 --- a/dynamicsizerecarray/version.py +++ b/dynamicsizerecarray/version.py @@ -1 +1 @@ -__version__ = "0.0.9" +__version__ = "0.0.10"