[mlir][OpDSL] Restructure comprehension.py (NFC).

Group and reorder the classed defined by comprehension.py and add type annotations.

Depends On D119126

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D119692