Driver was checking for direct mode and trying to lock it, but
left a gap where mode could change before the desired operation.
Use iio_device_claim_direct_mode() to guarantee device stays in
direct mode.

Refactor function to clarify look-up followed by lock sequence.

Signed-off-by: Alison Schofield <amsfield22@gmail.com>
Cc: Daniel Baluta <daniel.baluta@gmail.com>
Acked-by: Lars-Peter Clausen <lars@metafoo.de>
Signed-off-by: Jonathan Cameron <jic23@kernel.org>