Tab to next row in TreeTableView after last editable cell in one row

Question

I'm working with a JavaFx TreeTableView in Java. Most of my cells are editable, but some are not. When editing a cell, pressing tab jumps to the next editable cell - as expected, but when at the end of a row, it just exits editing. Is there a way to change this behavior, so that the next row is selected with its first editable cell?

Custom factories when adding a new column:

/**
 * Create a column. The column is also added to the billing table.
 * 
 * @param columnName
 *            Name of the column.
 * @return Created column.
 */
public TreeTableColumn<BillingTableRow, Double> addColumn(String columnName) {
    TreeTableColumn<BillingTableRow, Double> column = new TreeTableColumn<>(columnName);

    // Bind column data
    column.setCellValueFactory(new Callback<TreeTableColumn.CellDataFeatures<BillingTableRow, Double>, ObservableValue<Double>>() {
        @Override
        public ObservableValue<Double> call(CellDataFeatures<BillingTableRow, Double> param) {
            return new ObservableValueBase<Double>() {
                @Override
                public Double getValue() {
                    if (param.getValue().getValue().getType() == BillingTableRow.TYPE_SERVICE) {
                        if (param.getValue().getValue().getService().getAll().containsKey(columnName)) {
                            return param.getValue().getValue().getService().getAll().get(columnName).getTime();
                        }
                    }
                    return null;
                }

            };
        }
    });

    // Cell editing
    column.setCellFactory(new Callback<TreeTableColumn<BillingTableRow, Double>, TreeTableCell<BillingTableRow, Double>>() {
        @Override
        public TreeTableCell<BillingTableRow, Double> call(TreeTableColumn<BillingTableRow, Double> p) {
            return new EditableTreeTableDoubleCell() {
                @Override
                public void updateItem(Double item, boolean empty) {
                    super.updateItem(item, empty);
                    BillingTableRow row = getTreeTableRow().getItem();
                    if (!empty && row != null) {
                        // Styling
                        if (row.getType() == BillingTableRow.TYPE_CLIENT) {
                            setStyle("-fx-font-weight: bold;");
                        } else {
                            setStyle("");
                        }

                        // Update total values
                        BillingTableRow parentRow = (BillingTableRow) getTreeTableRow().getTreeItem().getParent().getValue();
                        parentRow.updateTotal();
                        row.getTable().getRoot().getValue().updateTotal();
                    } else {
                        // Clear cell if it's empty
                        setText(null);
                        setGraphic(null);
                    }
                    // if (empty) {
                    // setEditable(false);
                    // }
                }
            };
        }
    });

    // Edit event handler
    column.setOnEditCommit(new EventHandler<TreeTableColumn.CellEditEvent<BillingTableRow, Double>>() {
        @Override
        public void handle(CellEditEvent<BillingTableRow, Double> t) {
            BillingTableRow row = t.getTreeTablePosition().getTreeItem().getValue();
            if (t.getNewValue() != null) {
                if (row.getType() == BillingTableRow.TYPE_SERVICE && row.getService().getAll().containsKey(columnName)) {
                    row.getService().getAll().get(columnName).setTime(t.getNewValue());
                }
            }
            if (row.getChildren().size() > 0) {
                if (t.getTreeTablePosition().getTreeItem().getParent().getParent() != null)
                t.getTreeTablePosition().getTreeItem().getParent().getParent().getValue().updateTotal();
            }
        }
    });

    table.getColumns().add(column);
    return column;
}

Custom cells:

package gui;

import java.util.ArrayList;
import java.util.List;

import javafx.application.Platform;
import javafx.beans.value.ChangeListener;
import javafx.beans.value.ObservableValue;
import javafx.event.EventHandler;
import javafx.scene.control.ContentDisplay;
import javafx.scene.control.TextField;
import javafx.scene.control.TreeTableCell;
import javafx.scene.control.TreeTableColumn;
import javafx.scene.input.KeyCode;
import javafx.scene.input.KeyEvent;

/**
 * Editable TreeTableCell of type double.
 * 
 * @author Peter Jonsson, Graham Smith
 *
 */
public class EditableTreeTableDoubleCell extends TreeTableCell<BillingTableRow, Double> {
    private TextField textField;

    public EditableTreeTableDoubleCell() {
    }

    @Override
    public void startEdit() {
        super.startEdit();
        if (textField == null) {
            if (isEditable()) {
                createTextField();
            } else {
                return;
            }
        }
        setGraphic(textField);
        setContentDisplay(ContentDisplay.GRAPHIC_ONLY);
        Platform.runLater(new Runnable() {
            @Override
            public void run() {
                textField.requestFocus();
                textField.selectAll();
            }
        });
    }

    @Override
    public void cancelEdit() {
        super.cancelEdit();
        if (getItem() != null) {
            setText(Double.toString(getItem()));
        }
        setContentDisplay(ContentDisplay.TEXT_ONLY);
    }

    @Override
    public void updateItem(Double item, boolean empty) {
        super.updateItem(item, empty);
        if (empty) {
            setText(null);
            setGraphic(null);
        } else {
            if (isEditing()) {
                if (textField != null) {
                    textField.setText(getString());
                }
                setGraphic(textField);
                setContentDisplay(ContentDisplay.GRAPHIC_ONLY);
            } else {
                setText(getString());
                setContentDisplay(ContentDisplay.TEXT_ONLY);
            }
        }

        // Cell locking
        if (item == null) {
            setEditable(false);
        } else {
            setEditable(true);
        }
    }

    private void createTextField() {
        textField = new TextField(getString());
        if (textField.getText().isEmpty()) {
            textField.setText("0");
        }
        textField.setMinWidth(this.getWidth() - this.getGraphicTextGap() * 2);
        textField.setOnKeyPressed(new EventHandler<KeyEvent>() {
            @SuppressWarnings({ "unchecked", "rawtypes" })
            @Override
            public void handle(KeyEvent t) {
                if (t.getCode() == KeyCode.ENTER) {
                    commitEdit(Double.parseDouble(sanitizeInput(textField.getText())));

                } else if (t.getCode() == KeyCode.ESCAPE) {
                    cancelEdit();
                } else if (t.getCode() == KeyCode.TAB) {
                    commitEdit(Double.parseDouble(sanitizeInput(textField.getText())));
                    TreeTableColumn nextColumn = getNextColumn(!t.isShiftDown());
                    if (nextColumn != null) {
                        getTreeTableView().edit(getTreeTableRow().getIndex(), nextColumn);
                    }
                }
            }
        });
        textField.focusedProperty().addListener(new ChangeListener<Boolean>() {
            @Override
            public void changed(ObservableValue<? extends Boolean> observable, Boolean oldValue, Boolean newValue) {
                if (!newValue && textField != null) {
                    commitEdit(Double.parseDouble(sanitizeInput(textField.getText())));
                }
            }
        });
    }

    private String getString() {
        return getItem() == null ? "" : getItem().toString();
    }

    /**
     *
     * @param forward
     *            true gets the column to the right, false the column to the
     *            left of the current column
     * @return next column
     */
    private TreeTableColumn<BillingTableRow, ?> getNextColumn(boolean forward) {
        List<TreeTableColumn<BillingTableRow, ?>> columns = new ArrayList<>();
        for (TreeTableColumn<BillingTableRow, ?> column : getTreeTableView().getColumns()) {
            columns.addAll(getLeaves(column));
        }
        // There is no other column that supports editing.
        if (columns.size() < 2) {
            return null;
        }
        int currentIndex = columns.indexOf(getTableColumn());
        int nextIndex = currentIndex;
        if (forward) {
            nextIndex++;
            if (nextIndex > columns.size() - 1) {
                nextIndex = 0;
            }
        } else {
            nextIndex--;
            if (nextIndex < 0) {
                nextIndex = columns.size() - 1;
            }
        }
        return columns.get(nextIndex);
    }

    private List<TreeTableColumn<BillingTableRow, ?>> getLeaves(TreeTableColumn<BillingTableRow, ?> root) {
        List<TreeTableColumn<BillingTableRow, ?>> columns = new ArrayList<>();
        if (root.getColumns().isEmpty()) {
            // We only want the leaves that are editable.
            if (root.isEditable()) {
                columns.add(root);
            }
            return columns;
        } else {
            for (TreeTableColumn<BillingTableRow, ?> column : root.getColumns()) {
                columns.addAll(getLeaves(column));
            }
            return columns;
        }
    }

    /**
     * Sanitize an inputted string.
     * 
     * @param string
     *            String to sanitize.
     * @return Sanitized string.
     */
    private String sanitizeInput(String string) {
        if (string != null && !string.isEmpty()) {
            return string;
        }
        return "0";
    }
}