[datamodule] Slightly streamlines the datamodule approach; adds documentation; adds type annotations; adds TODOs