break some check from run_uvdata_uvsim into a separate function

Author: mkolopanis

pyuvsim/mpi.py

@@ -31,6 +31,8 @@
 
 
 class Tags(enum.IntEnum):
+    """Tags for MPI computations where worker nodes communicate with a main distribution node."""
+
     READY = enum.auto()
     START = enum.auto()
     DONE = enum.auto()

pyuvsim/utils.py

@@ -44,9 +44,11 @@ def __init__(self, maxval=None):
         self.remain = None
 
     def __enter__(self):
+        """Enter a context manager."""
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        """Exit a context manager."""
         self.finish()
 
     def update(self, count):

pyuvsim/uvsim.py

@@ -887,6 +887,41 @@ def _set_nsky_parts(Nsrcs, cat_nfreqs, Nsky_parts):
     return Nsky_parts
 
 
+def _check_uvdata(input_uv):
+    """Perform UVData checks for run_uvdata_uvsim.
+
+    Moved checks to reduce complexity.
+
+    Parameters
+    -----------
+    input_uv : :class:`pyuvdata.UVData` instance
+        Provides baseline/time/frequency information.
+
+    Returns
+    -------
+    the UVData ordering convention
+    """
+    if not isinstance(input_uv, UVData):
+        raise TypeError("input_uv must be UVData object")
+
+    if not (
+        (input_uv.Npols == 4)
+        and (input_uv.polarization_array.tolist() == [-5, -6, -7, -8])
+    ):
+        raise ValueError("input_uv must have XX,YY,XY,YX polarization")
+
+    if not input_uv.future_array_shapes:
+        input_uv.use_future_array_shapes()
+
+    input_order = input_uv.blt_order
+    if input_order != ("time", "baseline"):
+        input_uv.reorder_blts(order="time", minor_order="baseline")
+
+    return input_order
+
+
+# this function is very complex according to flake8.
+# too many ifs, we can swing back to this later.
 def run_uvdata_uvsim(
     input_uv,
     beam_list,
@@ -986,21 +1021,7 @@ def run_uvdata_uvsim(
     beam_dict = comm.bcast(beam_dict, root=0)
     catalog.share(root=0)
 
-    if not isinstance(input_uv, UVData):
-        raise TypeError("input_uv must be UVData object")
-
-    if not (
-        (input_uv.Npols == 4)
-        and (input_uv.polarization_array.tolist() == [-5, -6, -7, -8])
-    ):
-        raise ValueError("input_uv must have XX,YY,XY,YX polarization")
-
-    if not input_uv.future_array_shapes:
-        input_uv.use_future_array_shapes()
-
-    input_order = input_uv.blt_order
-    if input_order != ("time", "baseline"):
-        input_uv.reorder_blts(order="time", minor_order="baseline")
+    input_order = _check_uvdata(input_uv)
 
     # The root node will initialize our simulation
     # Read input file and make uvtask list