Skip to content

Commit db43bd1

Browse files
committed
Auto-expand string width for column insert
1 parent 22fe622 commit db43bd1

5 files changed

Lines changed: 55 additions & 10 deletions

File tree

‎CHANGES.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -737,6 +737,10 @@ astropy.table
737737
as long as the custom unit is enabled during reading (otherwise, the unit
738738
will become an ``UnrecognizedUnit``). [#9015]
739739

740+
- Fix bug where string values could be truncated when inserting into a
741+
``Column`` or ``MaskedColumn``, or when adding or inserting a row containing
742+
string values. [#9559]
743+
740744
astropy.tests
741745
^^^^^^^^^^^^^
742746

‎astropy/table/column.py‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,41 @@ def __setitem__(self, item, val):
125125
.format(self.__class__.__name__))
126126

127127

128+
def _expand_string_array_for_values(arr, values):
129+
"""
130+
For string-dtype return a version of ``arr`` that is wide enough for ``values``.
131+
If ``arr`` is not string-dtype or does not need expansion then return ``arr``.
132+
133+
Parameters
134+
----------
135+
arr : np.ndarray
136+
Input array
137+
values : scalar or array-like
138+
Values for width comparison for string arrays
139+
140+
Returns
141+
-------
142+
arr_expanded : np.ndarray
143+
144+
"""
145+
# If this is not a string type then just return original arr
146+
if arr.dtype.kind not in ('U', 'S'):
147+
return arr
148+
149+
# Find the length of the longest string in the new values.
150+
values_str_len = np.char.str_len(values).max()
151+
152+
# Determine character repeat count of self.dtype. Returns a positive
153+
# int or None (something like 'U0' is not possible in numpy). If new values
154+
# are longer than current then make a new (wider) version of self.
155+
arr_str_len = dtype_bytes_or_chars(arr.dtype)
156+
if arr_str_len and values_str_len > arr_str_len:
157+
arr_dtype = arr.dtype.byteorder + arr.dtype.kind + str(values_str_len)
158+
return arr.astype(arr_dtype)
159+
else:
160+
return arr
161+
162+
128163
class ColumnInfo(BaseColumnInfo):
129164
"""
130165
Container for meta information like name, description, format.
@@ -1030,10 +1065,9 @@ def insert(self, obj, values, axis=0):
10301065
data = np.insert(self, obj, None, axis=axis)
10311066
data[obj] = values
10321067
else:
1033-
# Explicitly convert to dtype of this column. Needed because numpy 1.7
1034-
# enforces safe casting by default, so . This isn't the case for 1.6 or 1.8+.
1035-
values = np.asarray(values, dtype=self.dtype)
1036-
data = np.insert(self, obj, values, axis=axis)
1068+
self_for_insert = _expand_string_array_for_values(self, values)
1069+
data = np.insert(self_for_insert, obj, values, axis=axis)
1070+
10371071
out = data.view(self.__class__)
10381072
out.__array_finalize__(self)
10391073
return out
@@ -1343,16 +1377,14 @@ def insert(self, obj, values, mask=None, axis=0):
13431377
new_data = np.insert(self_ma.data, obj, None, axis=axis)
13441378
new_data[obj] = values
13451379
else:
1346-
# Explicitly convert to dtype of this column. Needed because numpy 1.7
1347-
# enforces safe casting by default, so . This isn't the case for 1.6 or 1.8+.
1348-
values = np.asarray(values, dtype=self.dtype)
1380+
self_ma = _expand_string_array_for_values(self_ma, values)
13491381
new_data = np.insert(self_ma.data, obj, values, axis=axis)
13501382

13511383
if mask is None:
13521384
if self.dtype.kind == 'O':
13531385
mask = False
13541386
else:
1355-
mask = np.zeros(values.shape, dtype=bool)
1387+
mask = np.zeros(np.broadcast(values).shape, dtype=bool)
13561388
new_mask = np.insert(self_ma.mask, obj, mask, axis=axis)
13571389
new_ma = np.ma.array(new_data, mask=new_mask, copy=False)
13581390

‎astropy/table/tests/test_column.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,15 @@ def test_insert_axis(self, Column):
319319
c1 = c.insert(1, [5, 6], axis=1)
320320
assert np.all(c1 == [[1, 5, 2], [3, 6, 4]])
321321

322+
def test_insert_string_expand(self, Column):
323+
c = Column(['a', 'b'])
324+
c1 = c.insert(0, 'abc')
325+
assert np.all(c1 == ['abc', 'a', 'b'])
326+
327+
c = Column(['a', 'b'])
328+
c1 = c.insert(0, ['c', 'def'])
329+
assert np.all(c1 == ['c', 'def', 'a', 'b'])
330+
322331
def test_insert_multidim(self, Column):
323332
c = Column([[1, 2],
324333
[3, 4]], name='a', dtype=int)

‎astropy/table/tests/test_table.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -693,7 +693,7 @@ def test_add_with_list(self, table_types):
693693
assert len(t) == 4
694694
assert np.all(t['a'] == np.array([1, 2, 3, 4]))
695695
assert np.allclose(t['b'], np.array([4.0, 5.1, 6.2, 7.2]))
696-
assert np.all(t['c'] == np.array(['7', '8', '9', '1']))
696+
assert np.all(t['c'] == np.array(['7', '8', '9', '10']))
697697

698698
def test_add_with_dict(self, table_types):
699699
self._setup(table_types)

‎docs/table/index.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,7 @@ Adding a new row of data to the table is as follows. Note that the unit
228228
value is given in ``cm / s`` but will be added to the table as ``0.1 m / s`` in
229229
accord with the existing unit.
230230

231-
>>> t.add_row([-8, -9, 10 * u.cm / u.s, 11])
231+
>>> t.add_row([-8, 'string', 10 * u.cm / u.s, 10])
232232
>>> len(t)
233233
4
234234

0 commit comments

Comments
 (0)