Merge pull request #25 from tnakazato/refactor-filler-observation

refactor: filler for OBSERVATION table
tnakazato · Feb 3, 2025 · 5612935 · 5612935
2 parents 3a770b6 + 778de70
commit 5612935
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 29 deletions.
diff --git a/src/nro45data/psw/ms2/filler/__init__.py b/src/nro45data/psw/ms2/filler/__init__.py
@@ -9,7 +9,7 @@
 from .data_description import fill_data_description
 from .feed import fill_feed
 from .field import fill_field
-from .observation import _fill_observation_columns, _get_observation_columns
+from .observation import fill_observation
 from .polarization import _fill_polarization_columns, _get_polarization_columns
 from .main import fill_main
 from .pointing import fill_pointing
@@ -26,11 +26,6 @@
 LOG = logging.getLogger(__name__)
 
 
-def fill_observation(msfile: str, hdu: BinTableHDU):
-    columns = _get_observation_columns(hdu)
-    _fill_observation_columns(msfile, columns)
-
-
 def fill_polarization(msfile: str, hdu: BinTableHDU):
     columns = _get_polarization_columns(hdu)
     _fill_polarization_columns(msfile, columns)

diff --git a/src/nro45data/psw/ms2/filler/observation.py b/src/nro45data/psw/ms2/filler/observation.py
@@ -1,18 +1,28 @@
+from __future__ import annotations
+
 import logging
-from typing import TYPE_CHECKING
+from typing import Generator, TYPE_CHECKING
 
 import numpy as np
 
-from .._casa import open_table, datestr2mjd
-from .utils import fix_nrow_to
+from .._casa import datestr2mjd
+from .utils import fill_ms_table
 
 if TYPE_CHECKING:
     from astropy.io.fits.hdu.BinTableHDU import BinTableHDU
 
 LOG = logging.getLogger(__name__)
 
 
-def _get_observation_columns(hdu: "BinTableHDU") -> dict:
+def _get_observation_row(hdu: BinTableHDU) -> Generator[dict, None, None]:
+    """Provide observation row information.
+
+    Args:
+        hdu: NRO45m psw data in the form of BinTableHDU object.
+
+    Yields:
+        Dictionary containing observation row information.
+    """
     history_cards = hdu.header["HISTORY"]
 
     # TELESCOPE_NAME
@@ -45,17 +55,17 @@ def _get_observation_columns(hdu: "BinTableHDU") -> dict:
 
     # SCHEDULE_TYPE
     schedule_card = [x for x in history_cards if x.startswith("NEWSTAR SCHED")]
-    _schedule_type = schedule_card[0].split("=")[-1].strip(" '")
+    _schedule_type = schedule_card[0].split("=", maxsplit=1)[-1].strip(" '")
     schedule_type = f"{telescope_name} {_schedule_type}"
     LOG.debug("schedule_type: %s", schedule_type)
 
     # SCHEDULE
     group_card = [x for x in history_cards if x.startswith("NEWSTAR GROUP")]
-    _group = group_card[0].split("=")[-1].strip(" '")
+    _group = group_card[0].split("=", maxsplit=1)[-1].strip(" '")
     title1_card = [x for x in history_cards if x.startswith("NEWSTAR TITLE1")]
-    _title1 = title1_card[0].split("=")[-1].strip(" '")
+    _title1 = title1_card[0].split("=", maxsplit=1)[-1].strip(" '")
     title2_card = [x for x in history_cards if x.startswith("NEWSTAR TITLE2")]
-    _title2 = title2_card[0].split("=")[-1].strip(" '")
+    _title2 = title2_card[0].split("=", maxsplit=1)[-1].strip(" '")
     schedule = [f"{telescope_name} {_group}"]
     if _title1:
         schedule.append(_title1)
@@ -65,7 +75,7 @@ def _get_observation_columns(hdu: "BinTableHDU") -> dict:
 
     # PROJECT
     project_card = [x for x in history_cards if x.startswith("NEWSTAR PROJECT")]
-    project = project_card[0].split("=")[-1].strip(" '")
+    project = project_card[0].split("=", maxsplit=1)[-1].strip(" '")
     LOG.debug("project: %s", project)
 
     # RELEASE_DATE
@@ -74,7 +84,7 @@ def _get_observation_columns(hdu: "BinTableHDU") -> dict:
     # FLAG_ROW
     flag_row = False
 
-    columns = {
+    row = {
         "TELESCOPE_NAME": telescope_name,
         "TIME_RANGE": time_range,
         "OBSERVER": observer,
@@ -86,19 +96,14 @@ def _get_observation_columns(hdu: "BinTableHDU") -> dict:
         "FLAG_ROW": flag_row,
     }
 
-    return columns
+    yield row
 
 
-def _fill_observation_columns(msfile: str, columns: dict):
-    with open_table(msfile + "/OBSERVATION", read_only=False) as tb:
-        fix_nrow_to(1, tb)
+def fill_observation(msfile: str, hdu: BinTableHDU):
+    """Fill MS OBSERVATION table.
 
-        tb.putcell("TELESCOPE_NAME", 0, columns["TELESCOPE_NAME"])
-        tb.putcell("TIME_RANGE", 0, columns["TIME_RANGE"])
-        tb.putcell("OBSERVER", 0, columns["OBSERVER"])
-        tb.putcell("LOG", 0, columns["LOG"])
-        tb.putcell("SCHEDULE_TYPE", 0, columns["SCHEDULE_TYPE"])
-        tb.putcell("SCHEDULE", 0, columns["SCHEDULE"])
-        tb.putcell("PROJECT", 0, columns["PROJECT"])
-        tb.putcell("RELEASE_DATE", 0, columns["RELEASE_DATE"])
-        tb.putcell("FLAG_ROW", 0, columns["FLAG_ROW"])
+    Args:
+        msfile: Name of MS file.
+        hdu: NRO45m psw data in the form of BinTableHDU object.
+    """
+    fill_ms_table(msfile, hdu, "OBSERVATION", _get_observation_row)
diff --git a/tests/ms2_filler/test_forest.py b/tests/ms2_filler/test_forest.py
@@ -146,6 +146,19 @@ def test_forest_ms2_structure(msfile):
         end_time = mjd2datetime(time_range[1])
         # impose msec accuracy
         assert abs((end_time - end_expected).total_seconds()) < 1e-3
+        assert tb.getcell("TELESCOPE_NAME", 0) == "NRO45M"
+        assert tb.getcell("OBSERVER", 0) == "csv24"
+        log = tb.getcell("LOG", 0)
+        assert len(log) == 1
+        assert log[0] == ""
+        assert tb.getcell("SCHEDULE_TYPE", 0) == "NRO45M nmlfr"
+        schedule = tb.getcell("SCHEDULE", 0)
+        assert len(schedule) == 2
+        assert schedule[0] == "NRO45M csv24"
+        assert schedule[1] == "ffix                           z=1 3b= 0"
+        assert tb.getcell("PROJECT", 0) == "squint"
+        assert tb.getcell("RELEASE_DATE", 0) == 0
+        assert tb.getcell("FLAG_ROW", 0) is False
 
     with open_table(os.path.join(msfile, "SPECTRAL_WINDOW")) as tb:
         assert tb.nrows() == num_spws

diff --git a/tests/ms2_filler/test_h40.py b/tests/ms2_filler/test_h40.py
@@ -148,6 +148,19 @@ def test_h40_ms2_structure(msfile):
         end_time = mjd2datetime(time_range[1])
         # impose msec accuracy
         assert abs((end_time - end_expected).total_seconds()) < 1e-3
+        assert tb.getcell("TELESCOPE_NAME", 0) == "NRO45M"
+        assert tb.getcell("OBSERVER", 0) == "csv24"
+        log = tb.getcell("LOG", 0)
+        assert len(log) == 1
+        assert log[0] == ""
+        assert tb.getcell("SCHEDULE_TYPE", 0) == "NRO45M nmlh40"
+        schedule = tb.getcell("SCHEDULE", 0)
+        assert len(schedule) == 2
+        assert schedule[0] == "NRO45M csv24"
+        assert schedule[1] == "ffix                           z=1 3b= 0"
+        assert tb.getcell("PROJECT", 0) == "squint"
+        assert tb.getcell("RELEASE_DATE", 0) == 0
+        assert tb.getcell("FLAG_ROW", 0) is False
 
     with open_table(os.path.join(msfile, "SPECTRAL_WINDOW")) as tb:
         assert tb.nrows() == num_spws

diff --git a/tests/ms2_filler/test_z45.py b/tests/ms2_filler/test_z45.py
@@ -147,6 +147,19 @@ def test_z45_ms2_structure(msfile):
         end_time = mjd2datetime(time_range[1])
         # impose msec accuracy
         assert abs((end_time - end_expected).total_seconds()) < 1e-3
+        assert tb.getcell("TELESCOPE_NAME", 0) == "NRO45M"
+        assert tb.getcell("OBSERVER", 0) == "csv24"
+        log = tb.getcell("LOG", 0)
+        assert len(log) == 1
+        assert log[0] == ""
+        assert tb.getcell("SCHEDULE_TYPE", 0) == "NRO45M nmlzrf"
+        schedule = tb.getcell("SCHEDULE", 0)
+        assert len(schedule) == 2
+        assert schedule[0] == "NRO45M csv24"
+        assert schedule[1] == "ffix                           z=1 3b= 0"
+        assert tb.getcell("PROJECT", 0) == "squint"
+        assert tb.getcell("RELEASE_DATE", 0) == 0
+        assert tb.getcell("FLAG_ROW", 0) is False
 
     with open_table(os.path.join(msfile, "SPECTRAL_WINDOW")) as tb:
         assert tb.nrows() == num_spws