Repository: boto/boto3
File: boto3/dynamodb/table.py (Python)

Code before:

    def _extract_pkey_values(self, request):
        if request.get('PutRequest'):
            return [
                request['PutRequest']['Item'][key]
                for key in self._overwrite_by_pkeys
            ]
        elif request.get('DeleteRequest'):
            return [
                request['DeleteRequest']['Key'][key]
                for key in self._overwrite_by_pkeys
            ]
        return None

    def _flush_if_needed(self):
        if len(self._items_buffer) >= self._flush_amount:
            self._flush()

    def _flush(self):
        items_to_send = self._items_buffer[: self._flush_amount]
        self._items_buffer = self._items_buffer[self._flush_amount :]
        response = self._client.batch_write_item(
            RequestItems={self._table_name: items_to_send}
        )
        unprocessed_items = response['UnprocessedItems']
        if unprocessed_items is None:
            unprocessed_items = []
        else:
            unprocessed_items = unprocessed_items.get(self._table_name, [])
        # Any unprocessed_items are immediately added to the
        # next batch we send.
        self._items_buffer.extend(unprocessed_items)
        logger.debug(
            "Batch write sent %s, unprocessed: %s",
            len(items_to_send),
            len(self._items_buffer),
        )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        # When we exit, we need to keep flushing whatever's left
        # until there's nothing left in our items buffer.
        while self._items_buffer:
            self._flush()


Review comment (bug):
The semantics for this change are slightly different than what I'd proposed. This change accounts explicitly for `None`, but we were previously checking for anything falsey. For the sake of robustness the `if not unprocessed_items` may prevent any unintended failures.

Also it looks like you set `unprocessed_items` to an empty list. Ideally we want to be minimize type swapping for the same variable within a function. It obviously happens because we're accounting for falsey types above, but we should aim to avoid it. Otherwise, it's going to make type hints harder for us if they're added in the future. Was there an added benefit with this approach?

Code after:

    def _extract_pkey_values(self, request):
        if request.get('PutRequest'):
            return [
                request['PutRequest']['Item'][key]
                for key in self._overwrite_by_pkeys
            ]
        elif request.get('DeleteRequest'):
            return [
                request['DeleteRequest']['Key'][key]
                for key in self._overwrite_by_pkeys
            ]
        return None

    def _flush_if_needed(self):
        if len(self._items_buffer) >= self._flush_amount:
            self._flush()

    def _flush(self):
        items_to_send = self._items_buffer[: self._flush_amount]
        self._items_buffer = self._items_buffer[self._flush_amount :]
        response = self._client.batch_write_item(
            RequestItems={self._table_name: items_to_send}
        )
        unprocessed_items = response['UnprocessedItems']
        if not unprocessed_items:
            unprocessed_items = {}
        item_list = unprocessed_items.get(self._table_name, [])
        # Any unprocessed_items are immediately added to the
        # next batch we send.
        self._items_buffer.extend(item_list)
        logger.debug(
            "Batch write sent %s, unprocessed: %s",
            len(items_to_send),
            len(self._items_buffer),
        )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        # When we exit, we need to keep flushing whatever's left
        # until there's nothing left in our items buffer.
        while self._items_buffer:
            self._flush()