updating pytables values with table.itersorted()

itersorted_example.py

import tables


path = "./itersorted_example.h5"


class Value(tables.IsDescription):
    index = tables.IntCol()
    value = tables.IntCol()


def write_some_data():
    with tables.openFile(path, mode='w', title='test') as h5file:
        table = h5file.createTable('/', 'test', Value, 'Test Values')
        table.cols.index.createCSIndex()

        value_row = table.row

        for i in xrange(100):
            value_row['index'] = i
            value_row['value'] = 1
            value_row.append()
            table.flush()

    h5file.close()


def update_with_where():
    with tables.openFile(path, mode='r+') as h5file:
        table = h5file.root.test
        for i in xrange(10):
            where_clause = "index == " + str(i)
            for existing_row in table.where(where_clause):
                assert existing_row['index'] == i
                existing_row['value'] = 2
                existing_row.update()
                table.flush()
    h5file.close()


def update_with_itersorted():
    with tables.openFile(path, mode='r+') as h5file:
        table = h5file.root.test
        table_iterator = table.itersorted('index')
        for i in xrange(10):
            existing_row = table_iterator.next()
            assert existing_row['index'] == i
            existing_row['value'] = 3
            existing_row.update()
            table.flush()
    h5file.close()


def update_with_itersorted_workaround():
    with tables.openFile(path, mode='r+') as h5file:
        table = h5file.root.test
        table_iterator = table.itersorted('index')
        for i in xrange(10):
            existing_row = table_iterator.next()
            assert existing_row['index'] == i

            # update the row tuple directly on the table
            # note: this probably doesn't work if you change the column you are sorting by
            table[existing_row.nrow] = (i, 4)
            table.flush()
    h5file.close()


def assert_table_value(index, value):
    with tables.openFile(path, mode='r') as h5file:
        table = h5file.root.test
        where_clause = "index == " + str(index)
        # should only be one value, but a for loop handles the iteration nicely
        for table_value in table.where(where_clause):
            assert table_value['index'] == index
            assert table_value['value'] == value


if __name__ == '__main__':
    # first write out some data
    write_some_data()
    assert_table_value(5, 1)

    # then update first ten rows using table.where()
    update_with_where()
    assert_table_value(5, 2)

    # then update rows where the row comes from table.itersorted(); this fails
    update_with_itersorted()
    assert_table_value(5, 3)

    # this works if you comment out the preceding assert
    update_with_itersorted_workaround()
    assert_table_value(5, 4)